<commit_before>#[macro_use] extern crate custom_derive;

trait TryFrom<Src> {
    type Err;
    fn try_from(src: Src) -> Result<Self, Self::Err>;
}

macro_rules! TryFrom {
    (($prim:ty) $(pub)* enum $name:ident { $($body:tt)* }) => {
        TryFrom! {
            @collect_variants ($name, $prim),
            ($($body)*,) -> ()
        }
    };

    (
        @collect_variants ($name:ident, $prim:ty),
        ($(,)*) -> ($($var_names:ident,)*)
    ) => {
        impl TryFrom<$prim> for $name {
            type Err = $prim;
            fn try_from(src: $prim) -> Result<$name, $prim> {
                $(
                    if src == $name::$var_names as $prim {
                        return Ok($name::$var_names);
                    }
                )*
                Err(src)
            }
        }
    };

    (
        @collect_variants $fixed:tt,
        ($var:ident $(= $_val:expr)*, $($tail:tt)*) -> ($($var_names:tt)*)
    ) => {
        TryFrom! {
            @collect_variants $fixed,
            ($($tail)*) -> ($($var_names)* $var,)
        }
    };

    (
        @collect_variants ($name:ident),
        ($var:ident $_struct:tt, $($tail:tt)*) -> ($($var_names:tt)*)
    ) => {
        const _error: () = concat!(
            "cannot derive TryFrom for ",
            stringify!($name),
            ", due to non-unitary variant ",
            stringify!($var),
            "."
        );
    };
}

custom_derive! {
    #[derive(Debug, PartialEq, TryFrom(u8))]
    enum Get { Up, Down, AllAround }
}

#[test]
fn test_try_from() {
    assert_eq!(Get::try_from(0u8), Ok(Get::Up));
    assert_eq!(Get::try_from(1u8), Ok(Get::Down));
    assert_eq!(Get::try_from(2u8), Ok(Get::AllAround));
    assert_eq!(Get::try_from(3u8), Err(3u8));
}
<commit_msg>Fix for RFC 1214.<commit_after>#[macro_use] extern crate custom_derive;

trait TryFrom<Src>: Sized {
    type Err;
    fn try_from(src: Src) -> Result<Self, Self::Err>;
}

macro_rules! TryFrom {
    (($prim:ty) $(pub)* enum $name:ident { $($body:tt)* }) => {
        TryFrom! {
            @collect_variants ($name, $prim),
            ($($body)*,) -> ()
        }
    };

    (
        @collect_variants ($name:ident, $prim:ty),
        ($(,)*) -> ($($var_names:ident,)*)
    ) => {
        impl TryFrom<$prim> for $name {
            type Err = $prim;
            fn try_from(src: $prim) -> Result<$name, $prim> {
                $(
                    if src == $name::$var_names as $prim {
                        return Ok($name::$var_names);
                    }
                )*
                Err(src)
            }
        }
    };

    (
        @collect_variants $fixed:tt,
        ($var:ident $(= $_val:expr)*, $($tail:tt)*) -> ($($var_names:tt)*)
    ) => {
        TryFrom! {
            @collect_variants $fixed,
            ($($tail)*) -> ($($var_names)* $var,)
        }
    };

    (
        @collect_variants ($name:ident),
        ($var:ident $_struct:tt, $($tail:tt)*) -> ($($var_names:tt)*)
    ) => {
        const _error: () = concat!(
            "cannot derive TryFrom for ",
            stringify!($name),
            ", due to non-unitary variant ",
            stringify!($var),
            "."
        );
    };
}

custom_derive! {
    #[derive(Debug, PartialEq, TryFrom(u8))]
    enum Get { Up, Down, AllAround }
}

#[test]
fn test_try_from() {
    assert_eq!(Get::try_from(0u8), Ok(Get::Up));
    assert_eq!(Get::try_from(1u8), Ok(Get::Down));
    assert_eq!(Get::try_from(2u8), Ok(Get::AllAround));
    assert_eq!(Get::try_from(3u8), Err(3u8));
}
<|endoftext|>


<commit_before>use std::mem;

use Result;
use band::Band;
use primitive::*;

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum NamingTable {
    Format0(NamingTable0),
    Format1(NamingTable1),
}

spec! {
    pub NamingTable0 {
        format       (USHORT         ),
        count        (USHORT         ),
        stringOffset (USHORT         ),
        nameRecord   (Vec<NameRecord>) |band, this| { read_vector!(band, this.count) },
        storage      (Vec<BYTE>      ) |band, this| { this.read_storage(band) },
    }
}

spec! {
    pub NamingTable1 {
        format        (USHORT                ),
        count         (USHORT                ),
        stringOffset  (USHORT                ),
        nameRecord    (Vec<NameRecord>       ) |band, this| { read_vector!(band, this.count) },
        langTagCount  (USHORT                ),
        langTagRecord (Vec<LanguageTagRecord>) |band, this| { read_vector!(band,
                                                                           this.langTagCount) },
        storage       (Vec<BYTE>             ) |band, this| { this.read_storage(band) },
    }
}

spec! {
    #[repr(C)]
    #[derive(Copy)]
    pub NameRecord {
        platformID (USHORT),
        encodingID (USHORT),
        languageID (USHORT),
        nameID     (USHORT),
        length     (USHORT),
        offset     (USHORT),
    }
}

spec! {
    #[repr(C)]
    #[derive(Copy)]
    pub LanguageTagRecord {
        length (USHORT),
        ffset  (USHORT),
    }
}

impl NamingTable0 {
    #[inline]
    pub fn strings(&self) -> Result<Vec<String>> {
        strings(&self.nameRecord, &self.storage)
    }

    fn read_storage<T: Band>(&self, band: &mut T) -> Result<Vec<u8>> {
        let current = try!(band.position());
        let above = 3 * mem::size_of::<USHORT>() +
                    self.nameRecord.len() * mem::size_of::<NameRecord>();
        try!(band.jump(current - above as u64 + self.stringOffset as u64));
        read_vector!(band, storage_length(&self.nameRecord))
    }
}

impl NamingTable1 {
    #[inline]
    pub fn strings(&self) -> Result<Vec<String>> {
        strings(&self.nameRecord, &self.storage)
    }

    fn read_storage<T: Band>(&self, band: &mut T) -> Result<Vec<u8>> {
        let current = try!(band.position());
        let above = 4 * mem::size_of::<USHORT>() +
                    self.nameRecord.len() * mem::size_of::<NameRecord>() +
                    self.langTagRecord.len() * mem::size_of::<LanguageTagRecord>();
        try!(band.jump(current - above as u64 + self.stringOffset as u64));
        read_vector!(band, storage_length(&self.nameRecord))
    }
}

fn storage_length(records: &[NameRecord]) -> usize {
    let mut length = 0;
    for record in records {
        let end = record.offset + record.length + 1;
        if end > length {
            length = end;
        }
    }
    length as usize
}

fn strings(records: &[NameRecord], storage: &[u8]) -> Result<Vec<String>> {
    let mut strings = vec![];
    for record in records {
        let (offset, length) = (record.offset as usize, record.length as usize);
        let bytes = &storage[offset..(offset + length)];
        match record.platformID {
            1 => match decodeMacintosh(bytes, record.encodingID) {
                Some(string) => {
                    strings.push(string);
                    continue;
                },
                _ => {},
            },
            _ => {},
        }
        strings.push("<unsupported>".to_string());
    }
    Ok(strings)
}

// The implementation is based on
// https://github.com/nodebox/opentype.js/blob/master/src/types.js#L300
fn decodeMacintosh(bytes: &[BYTE], encoding: USHORT) -> Option<String> {
    const ROMAN: [char; 128] = ['Ä', 'Å', 'Ç', 'É', 'Ñ', 'Ö', 'Ü', 'á', 'à', 'â', 'ä', 'ã', 'å',
                                'ç', 'é', 'è', 'ê', 'ë', 'í', 'ì', 'î', 'ï', 'ñ', 'ó', 'ò', 'ô',
                                'ö', 'õ', 'ú', 'ù', 'û', 'ü', '†', '°', '¢', '£', '§', '•', '¶',
                                'ß', '®', '©', '™', '´', '¨', '≠', 'Æ', 'Ø', '∞', '±', '≤', '≥',
                                '¥', 'µ', '∂', '∑', '∏', 'π', '∫', 'ª', 'º', 'Ω', 'æ', 'ø', '¿',
                                '¡', '¬', '√', 'ƒ', '≈', '∆', '«', '»', '…', ' ', 'À', 'Ã', 'Õ',
                                'Œ', 'œ', '–', '—', '“', '”', '‘', '’', '÷', '◊', 'ÿ', 'Ÿ', '⁄',
                                '€', '‹', '›', 'ﬁ', 'ﬂ', '‡', '·', '‚', '„', '‰', 'Â', 'Ê', 'Á',
                                'Ë', 'È', 'Í', 'Î', 'Ï', 'Ì', 'Ó', 'Ô', '', 'Ò', 'Ú', 'Û', 'Ù',
                                'ı', 'ˆ', '˜', '¯', '˘', '˙', '˚', '¸', '˝', '˛', 'ˇ'];

    if encoding != 0 {
        return None;
    }

    let table = &ROMAN;

    let mut string = String::new();
    for &byte in bytes {
        if byte <= 0x7F {
            string.push(byte as char);
        } else {
            string.push(table[(byte & 0x7F) as usize]);
        }
    }

    Some(string)
}
<commit_msg>Make a stylistic adjustment<commit_after>use std::mem;

use Result;
use band::Band;
use primitive::*;

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum NamingTable {
    Format0(NamingTable0),
    Format1(NamingTable1),
}

spec! {
    pub NamingTable0 {
        format       (USHORT         ),
        count        (USHORT         ),
        stringOffset (USHORT         ),
        nameRecord   (Vec<NameRecord>) |band, this| { read_vector!(band, this.count) },
        storage      (Vec<BYTE>      ) |band, this| { this.read_storage(band) },
    }
}

spec! {
    pub NamingTable1 {
        format        (USHORT                ),
        count         (USHORT                ),
        stringOffset  (USHORT                ),
        nameRecord    (Vec<NameRecord>       ) |band, this| { read_vector!(band, this.count) },
        langTagCount  (USHORT                ),
        langTagRecord (Vec<LanguageTagRecord>) |band, this| { read_vector!(band,
                                                                           this.langTagCount) },
        storage       (Vec<BYTE>             ) |band, this| { this.read_storage(band) },
    }
}

spec! {
    #[repr(C)]
    #[derive(Copy)]
    pub NameRecord {
        platformID (USHORT),
        encodingID (USHORT),
        languageID (USHORT),
        nameID     (USHORT),
        length     (USHORT),
        offset     (USHORT),
    }
}

spec! {
    #[repr(C)]
    #[derive(Copy)]
    pub LanguageTagRecord {
        length (USHORT),
        ffset  (USHORT),
    }
}

impl NamingTable0 {
    #[inline]
    pub fn strings(&self) -> Result<Vec<String>> {
        strings(&self.nameRecord, &self.storage)
    }

    fn read_storage<T: Band>(&self, band: &mut T) -> Result<Vec<u8>> {
        let current = try!(band.position());
        let above = 3 * mem::size_of::<USHORT>() +
                    self.nameRecord.len() * mem::size_of::<NameRecord>();
        try!(band.jump(current - above as u64 + self.stringOffset as u64));
        read_vector!(band, storage_length(&self.nameRecord))
    }
}

impl NamingTable1 {
    #[inline]
    pub fn strings(&self) -> Result<Vec<String>> {
        strings(&self.nameRecord, &self.storage)
    }

    fn read_storage<T: Band>(&self, band: &mut T) -> Result<Vec<u8>> {
        let current = try!(band.position());
        let above = 4 * mem::size_of::<USHORT>() +
                    self.nameRecord.len() * mem::size_of::<NameRecord>() +
                    self.langTagRecord.len() * mem::size_of::<LanguageTagRecord>();
        try!(band.jump(current - above as u64 + self.stringOffset as u64));
        read_vector!(band, storage_length(&self.nameRecord))
    }
}

fn storage_length(records: &[NameRecord]) -> usize {
    let mut length = 0;
    for record in records {
        let end = record.offset + record.length + 1;
        if end > length {
            length = end;
        }
    }
    length as usize
}

fn strings(records: &[NameRecord], storage: &[u8]) -> Result<Vec<String>> {
    let mut strings = vec![];
    for record in records {
        let (offset, length) = (record.offset as usize, record.length as usize);
        let bytes = &storage[offset..(offset + length)];
        match record.platformID {
            1 => match decode_macintosh(bytes, record.encodingID) {
                Some(string) => {
                    strings.push(string);
                    continue;
                },
                _ => {},
            },
            _ => {},
        }
        strings.push("<unsupported>".to_string());
    }
    Ok(strings)
}

// The implementation is based on
// https://github.com/nodebox/opentype.js/blob/master/src/types.js#L300
fn decode_macintosh(bytes: &[BYTE], encoding: USHORT) -> Option<String> {
    const ROMAN: [char; 128] = ['Ä', 'Å', 'Ç', 'É', 'Ñ', 'Ö', 'Ü', 'á', 'à', 'â', 'ä', 'ã', 'å',
                                'ç', 'é', 'è', 'ê', 'ë', 'í', 'ì', 'î', 'ï', 'ñ', 'ó', 'ò', 'ô',
                                'ö', 'õ', 'ú', 'ù', 'û', 'ü', '†', '°', '¢', '£', '§', '•', '¶',
                                'ß', '®', '©', '™', '´', '¨', '≠', 'Æ', 'Ø', '∞', '±', '≤', '≥',
                                '¥', 'µ', '∂', '∑', '∏', 'π', '∫', 'ª', 'º', 'Ω', 'æ', 'ø', '¿',
                                '¡', '¬', '√', 'ƒ', '≈', '∆', '«', '»', '…', ' ', 'À', 'Ã', 'Õ',
                                'Œ', 'œ', '–', '—', '“', '”', '‘', '’', '÷', '◊', 'ÿ', 'Ÿ', '⁄',
                                '€', '‹', '›', 'ﬁ', 'ﬂ', '‡', '·', '‚', '„', '‰', 'Â', 'Ê', 'Á',
                                'Ë', 'È', 'Í', 'Î', 'Ï', 'Ì', 'Ó', 'Ô', '', 'Ò', 'Ú', 'Û', 'Ù',
                                'ı', 'ˆ', '˜', '¯', '˘', '˙', '˚', '¸', '˝', '˛', 'ˇ'];

    if encoding != 0 {
        return None;
    }

    let table = &ROMAN;

    let mut string = String::new();
    for &byte in bytes {
        if byte <= 0x7F {
            string.push(byte as char);
        } else {
            string.push(table[(byte & 0x7F) as usize]);
        }
    }

    Some(string)
}
<|endoftext|>


<commit_before>//! Disassembly calling function for most targets.

use crate::Function;
use std::{collections::HashSet, env, process::Command, str};

// Extracts the "shim" name from the `symbol`.
fn normalize(mut symbol: &str) -> String {
    // Remove trailing colon:
    if symbol.ends_with(':') {
        symbol = &symbol[..symbol.len() - 1];
    }
    if symbol.ends_with('>') {
        symbol = &symbol[..symbol.len() - 1];
    }
    if let Some(idx) = symbol.find('<') {
        symbol = &symbol[idx + 1..];
    }

    let mut symbol = rustc_demangle::demangle(symbol).to_string();
    symbol = match symbol.rfind("::h") {
        Some(i) => symbol[..i].to_string(),
        None => symbol.to_string(),
    };

    // Remove Rust paths
    if let Some(last_colon) = symbol.rfind(':') {
        symbol = (&symbol[last_colon + 1..]).to_string();
    }

    // Normalize to no leading underscore to handle platforms that may
    // inject extra ones in symbol names.
    while symbol.starts_with('_') {
        symbol.remove(0);
    }
    // Windows/x86 has a suffix such as @@4.
    if let Some(idx) = symbol.find("@@") {
        symbol = (&symbol[..idx]).to_string();
    }
    symbol
}

pub(crate) fn disassemble_myself() -> HashSet<Function> {
    let me = env::current_exe().expect("failed to get current exe");

    let disassembly = if cfg!(target_os = "windows") && cfg!(target_env = "msvc") {
        let target = if cfg!(target_arch = "x86_64") {
            "x86_64-pc-windows-msvc"
        } else if cfg!(target_arch = "x86") {
            "i686-pc-windows-msvc"
        } else {
            panic!("disassembly unimplemented")
        };
        let mut cmd = cc::windows_registry::find(target, "dumpbin.exe")
            .expect("failed to find `dumpbin` tool");
        let output = cmd
            .arg("/DISASM")
            .arg(&me)
            .output()
            .expect("failed to execute dumpbin");
        println!(
            "{}\n{}",
            output.status,
            String::from_utf8_lossy(&output.stderr)
        );
        assert!(output.status.success());
        // Windows does not return valid UTF-8 output:
        String::from_utf8_lossy(Vec::leak(output.stdout))
    } else if cfg!(target_os = "windows") {
        panic!("disassembly unimplemented")
    } else if cfg!(target_os = "macos") {
        let output = Command::new("otool")
            .arg("-vt")
            .arg(&me)
            .output()
            .expect("failed to execute otool");
        println!(
            "{}\n{}",
            output.status,
            String::from_utf8_lossy(&output.stderr)
        );
        assert!(output.status.success());

        String::from_utf8_lossy(Vec::leak(output.stdout))
    } else {
        let objdump = env::var("OBJDUMP").unwrap_or_else(|_| "objdump".to_string());
        let output = Command::new(objdump.clone())
            .arg("--disassemble")
            .arg(&me)
            .output()
            .unwrap_or_else(|_| panic!("failed to execute objdump. OBJDUMP={}", objdump));
        println!(
            "{}\n{}",
            output.status,
            String::from_utf8_lossy(&output.stderr)
        );
        assert!(output.status.success());

        String::from_utf8_lossy(Vec::leak(output.stdout))
    };

    parse(&disassembly)
}

fn parse(output: &str) -> HashSet<Function> {
    let mut lines = output.lines();

    println!(
        "First 100 lines of the disassembly input containing {} lines:",
        lines.clone().count()
    );
    for line in output.lines().take(100) {
        println!("{}", line);
    }

    let mut functions = HashSet::new();
    let mut cached_header = None;
    while let Some(header) = cached_header.take().or_else(|| lines.next()) {
        if !header.ends_with(':') || !header.contains("stdarch_test_shim") {
            continue;
        }
        eprintln!("header: {}", header);
        let symbol = normalize(header);
        eprintln!("normalized symbol: {}", symbol);
        let mut instructions = Vec::new();
        while let Some(instruction) = lines.next() {
            if instruction.ends_with(':') {
                cached_header = Some(instruction);
                break;
            }
            if instruction.is_empty() {
                cached_header = None;
                break;
            }
            let parts = if cfg!(target_os = "macos") {
                // Each line of instructions should look like:
                //
                //      $addr    $instruction...
                instruction
                    .split_whitespace()
                    .skip(1)
                    .map(std::string::ToString::to_string)
                    .collect::<Vec<String>>()
            } else if cfg!(target_env = "msvc") {
                // Each line looks like:
                //
                // >  $addr: ab cd ef     $instr..
                // >         00 12          # this line os optional
                if instruction.starts_with("       ") {
                    continue;
                }
                instruction
                    .split_whitespace()
                    .skip(1)
                    .skip_while(|s| s.len() == 2 && usize::from_str_radix(s, 16).is_ok())
                    .map(std::string::ToString::to_string)
                    .skip_while(|s| *s == "lock") // skip x86-specific prefix
                    .collect::<Vec<String>>()
            } else {
                // objdump
                // Each line of instructions should look like:
                //
                //      $rel_offset: ab cd ef 00    $instruction...
                let expected_len = if cfg!(target_arch = "arm") || cfg!(target_arch = "aarch64") {
                    8
                } else {
                    2
                };

                instruction
                    .split_whitespace()
                    .skip(1)
                    .skip_while(|s| s.len() == expected_len && usize::from_str_radix(s, 16).is_ok())
                    .skip_while(|s| *s == "lock") // skip x86-specific prefix
                    .map(std::string::ToString::to_string)
                    .collect::<Vec<String>>()
            };
            instructions.push(parts.join(" "));
        }
        let function = Function {
            name: symbol,
            instrs: instructions,
        };
        assert!(functions.insert(function));
    }

    eprintln!("all found functions dump:");
    for k in &functions {
        eprintln!("  f: {}", k.name);
    }

    functions
}
<commit_msg>Use --no-show-raw-insn to make disassemble parser simple. (#948)<commit_after>//! Disassembly calling function for most targets.

use crate::Function;
use std::{collections::HashSet, env, process::Command, str};

// Extracts the "shim" name from the `symbol`.
fn normalize(mut symbol: &str) -> String {
    // Remove trailing colon:
    if symbol.ends_with(':') {
        symbol = &symbol[..symbol.len() - 1];
    }
    if symbol.ends_with('>') {
        symbol = &symbol[..symbol.len() - 1];
    }
    if let Some(idx) = symbol.find('<') {
        symbol = &symbol[idx + 1..];
    }

    let mut symbol = rustc_demangle::demangle(symbol).to_string();
    symbol = match symbol.rfind("::h") {
        Some(i) => symbol[..i].to_string(),
        None => symbol.to_string(),
    };

    // Remove Rust paths
    if let Some(last_colon) = symbol.rfind(':') {
        symbol = (&symbol[last_colon + 1..]).to_string();
    }

    // Normalize to no leading underscore to handle platforms that may
    // inject extra ones in symbol names.
    while symbol.starts_with('_') {
        symbol.remove(0);
    }
    // Windows/x86 has a suffix such as @@4.
    if let Some(idx) = symbol.find("@@") {
        symbol = (&symbol[..idx]).to_string();
    }
    symbol
}

pub(crate) fn disassemble_myself() -> HashSet<Function> {
    let me = env::current_exe().expect("failed to get current exe");

    let disassembly = if cfg!(target_os = "windows") && cfg!(target_env = "msvc") {
        let target = if cfg!(target_arch = "x86_64") {
            "x86_64-pc-windows-msvc"
        } else if cfg!(target_arch = "x86") {
            "i686-pc-windows-msvc"
        } else {
            panic!("disassembly unimplemented")
        };
        let mut cmd = cc::windows_registry::find(target, "dumpbin.exe")
            .expect("failed to find `dumpbin` tool");
        let output = cmd
            .arg("/DISASM")
            .arg(&me)
            .output()
            .expect("failed to execute dumpbin");
        println!(
            "{}\n{}",
            output.status,
            String::from_utf8_lossy(&output.stderr)
        );
        assert!(output.status.success());
        // Windows does not return valid UTF-8 output:
        String::from_utf8_lossy(Vec::leak(output.stdout))
    } else if cfg!(target_os = "windows") {
        panic!("disassembly unimplemented")
    } else if cfg!(target_os = "macos") {
        let output = Command::new("otool")
            .arg("-vt")
            .arg(&me)
            .output()
            .expect("failed to execute otool");
        println!(
            "{}\n{}",
            output.status,
            String::from_utf8_lossy(&output.stderr)
        );
        assert!(output.status.success());

        String::from_utf8_lossy(Vec::leak(output.stdout))
    } else {
        let objdump = env::var("OBJDUMP").unwrap_or_else(|_| "objdump".to_string());
        let output = Command::new(objdump.clone())
            .arg("--disassemble")
            .arg("--no-show-raw-insn")
            .arg(&me)
            .output()
            .unwrap_or_else(|_| panic!("failed to execute objdump. OBJDUMP={}", objdump));
        println!(
            "{}\n{}",
            output.status,
            String::from_utf8_lossy(&output.stderr)
        );
        assert!(output.status.success());

        String::from_utf8_lossy(Vec::leak(output.stdout))
    };

    parse(&disassembly)
}

fn parse(output: &str) -> HashSet<Function> {
    let mut lines = output.lines();

    println!(
        "First 100 lines of the disassembly input containing {} lines:",
        lines.clone().count()
    );
    for line in output.lines().take(100) {
        println!("{}", line);
    }

    let mut functions = HashSet::new();
    let mut cached_header = None;
    while let Some(header) = cached_header.take().or_else(|| lines.next()) {
        if !header.ends_with(':') || !header.contains("stdarch_test_shim") {
            continue;
        }
        eprintln!("header: {}", header);
        let symbol = normalize(header);
        eprintln!("normalized symbol: {}", symbol);
        let mut instructions = Vec::new();
        while let Some(instruction) = lines.next() {
            if instruction.ends_with(':') {
                cached_header = Some(instruction);
                break;
            }
            if instruction.is_empty() {
                cached_header = None;
                break;
            }
            let parts = if cfg!(target_os = "macos") {
                // Each line of instructions should look like:
                //
                //      $addr    $instruction...
                instruction
                    .split_whitespace()
                    .skip(1)
                    .map(std::string::ToString::to_string)
                    .collect::<Vec<String>>()
            } else if cfg!(target_env = "msvc") {
                // Each line looks like:
                //
                // >  $addr: ab cd ef     $instr..
                // >         00 12          # this line os optional
                if instruction.starts_with("       ") {
                    continue;
                }
                instruction
                    .split_whitespace()
                    .skip(1)
                    .skip_while(|s| s.len() == 2 && usize::from_str_radix(s, 16).is_ok())
                    .map(std::string::ToString::to_string)
                    .skip_while(|s| *s == "lock") // skip x86-specific prefix
                    .collect::<Vec<String>>()
            } else {
                // objdump with --no-show-raw-insn
                // Each line of instructions should look like:
                //
                //      $rel_offset:       $instruction...
                instruction
                    .split_whitespace()
                    .skip(1)
                    .skip_while(|s| *s == "lock") // skip x86-specific prefix
                    .map(std::string::ToString::to_string)
                    .collect::<Vec<String>>()
            };
            instructions.push(parts.join(" "));
        }
        let function = Function {
            name: symbol,
            instrs: instructions,
        };
        assert!(functions.insert(function));
    }

    eprintln!("all found functions dump:");
    for k in &functions {
        eprintln!("  f: {}", k.name);
    }

    functions
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer, channels::Channels, connection_status::ConnectionState, frames::Frames,
    internal_rpc::InternalRPC, thread::ThreadHandle, waker::Waker, Configuration, ConnectionStatus,
    Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker as MioWaker};
use std::{
    collections::VecDeque,
    io::{Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

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

pub struct IoLoop<T> {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    connection_io_loop_handle: ThreadHandle,
    waker: Waker,
    socket: T,
    status: Status,
    poll: Poll,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
    sent_initial_heartbeat: bool,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        connection_io_loop_handle: ThreadHandle,
        waker: Waker,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        waker.set_waker(MioWaker::new(poll.registry(), WAKER)?);
        let frame_size = std::cmp::max(8192, configuration.frame_max() as usize);
        let mut inner = Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            connection_io_loop_handle,
            waker,
            socket,
            status: Status::Initial,
            poll,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
            sent_initial_heartbeat: false,
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat * 500)); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status == Status::Initial
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        self.run(&mut events)?;
                    }
                    Ok(())
                })?,
        );
        waker.wake()
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()?;
        self.internal_rpc.poll(&self.channels)
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            self.internal_rpc.poll(&self.channels)?;
            if self.connection_status.closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.internal_rpc.poll(&self.channels)?;
            if self.should_heartbeat() {
                self.sent_initial_heartbeat = true;
                self.channels.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection_status.errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection_status.state() {
            resolver.swear(Err(error.clone()));
            self.status = Status::Stop;
        }
        self.channels.set_connection_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            // If we didn't send the first heartbeat, don't update last_write to trick us into sending it
            if self.sent_initial_heartbeat {
                self.last_write = Instant::now();
            }

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

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                self.send_continue()?;
            }
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {:?}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            self.waker.wake()
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.channels.set_connection_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

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

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.channels.set_connection_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<commit_msg>IoLoop: ensure we poll internal_rpc on error<commit_after>use crate::{
    buffer::Buffer, channels::Channels, connection_status::ConnectionState, frames::Frames,
    internal_rpc::InternalRPC, thread::ThreadHandle, waker::Waker, Configuration, ConnectionStatus,
    Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker as MioWaker};
use std::{
    collections::VecDeque,
    io::{Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

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

pub struct IoLoop<T> {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    connection_io_loop_handle: ThreadHandle,
    waker: Waker,
    socket: T,
    status: Status,
    poll: Poll,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
    sent_initial_heartbeat: bool,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        connection_io_loop_handle: ThreadHandle,
        waker: Waker,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        waker.set_waker(MioWaker::new(poll.registry(), WAKER)?);
        let frame_size = std::cmp::max(8192, configuration.frame_max() as usize);
        let mut inner = Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            connection_io_loop_handle,
            waker,
            socket,
            status: Status::Initial,
            poll,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
            sent_initial_heartbeat: false,
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat * 500)); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status == Status::Initial
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        self.run(&mut events)?;
                    }
                    Ok(())
                })?,
        );
        waker.wake()
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        let res = self.do_run();
        self.internal_rpc.poll(&self.channels).and(res)
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            self.internal_rpc.poll(&self.channels)?;
            if self.connection_status.closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.internal_rpc.poll(&self.channels)?;
            if self.should_heartbeat() {
                self.sent_initial_heartbeat = true;
                self.channels.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection_status.errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection_status.state() {
            resolver.swear(Err(error.clone()));
            self.status = Status::Stop;
        }
        self.channels.set_connection_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            // If we didn't send the first heartbeat, don't update last_write to trick us into sending it
            if self.sent_initial_heartbeat {
                self.last_write = Instant::now();
            }

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

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                self.send_continue()?;
            }
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {:?}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            self.waker.wake()
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.channels.set_connection_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

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

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.channels.set_connection_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use llvm_sys::prelude::*;
use llvm_sys::core::*;

use libc::{c_int, c_uint};

use std::marker::PhantomData;
use std::mem::transmute_copy;

use id::IdRef;
use opaque::Opaque;
use inheritance::{upcast, DerivesFrom};

use llvm::value::Value;
use llvm::context::Context;

pub struct Type<'cid> {
    _context_id: IdRef<'cid>,
    _opaque: Opaque
}
unsafe impl<'cid> DerivesFrom<Type<'cid>> for Type<'cid> { }

pub struct PointerType<'cid, SubType: ?Sized> {
    _inner: PhantomData<SubType>,
    _super: Type<'cid>
}
unsafe impl<'cid, SubTypeSpecific: DerivesFrom<SubTypeGeneral> + ?Sized, SubTypeGeneral: ?Sized> DerivesFrom<PointerType<'cid, SubTypeGeneral>> for PointerType<'cid, SubTypeSpecific> { }
unsafe impl<'cid, General: ?Sized, SubType: ?Sized> DerivesFrom<General> for PointerType<'cid, SubType> where Type<'cid>: DerivesFrom<General> { }

pub struct ArrayType<'cid, SubType: ?Sized> {
    _inner: PhantomData<SubType>,
    _super: Type<'cid>
}
unsafe impl<'cid, SubTypeSpecific: DerivesFrom<SubTypeGeneral> + ?Sized, SubTypeGeneral: ?Sized> DerivesFrom<ArrayType<'cid, SubTypeGeneral>> for ArrayType<'cid, SubTypeSpecific> { }
unsafe impl<'cid, General: ?Sized, SubType: ?Sized> DerivesFrom<General> for ArrayType<'cid, SubType> where Type<'cid>: DerivesFrom<General> { }

pub struct FunctionType<'cid> {
    _super: Type<'cid>
}
unsafe impl<'cid> DerivesFrom<FunctionType<'cid>> for FunctionType<'cid> { }
unsafe impl<'cid, General: ?Sized> DerivesFrom<General> for FunctionType<'cid> where Type<'cid>: DerivesFrom<General> { }

pub struct IntegerType<'cid> {
    _super: Type<'cid>
}
unsafe impl<'cid> DerivesFrom<IntegerType<'cid>> for IntegerType<'cid> { }
unsafe impl<'cid, General: ?Sized> DerivesFrom<General> for IntegerType<'cid> where Type<'cid>: DerivesFrom<General> { }

impl<'cid> Type<'cid> {
    pub fn dump(&self) {
        unsafe {
            LLVMDumpType(self.as_raw())
        }
    }

    pub fn as_raw(&self) -> LLVMTypeRef {
        self as *const Type as *mut Type as LLVMTypeRef
    }

    pub fn f32<'ctx>(context: &'ctx Context<'cid>) -> &'ctx Type<'cid> {
        unsafe {
            &*(LLVMFloatTypeInContext(context.as_raw()) as *mut Type)
        }
    }

    pub fn f64<'ctx>(context: &'ctx Context<'cid>) -> &'ctx Type<'cid> {
        unsafe {
            &*(LLVMDoubleTypeInContext(context.as_raw()) as *mut Type)
        }
    }

    pub fn void<'ctx>(context: &'ctx Context<'cid>) -> &'ctx Type<'cid> {
        unsafe {
            &*(LLVMVoidTypeInContext(context.as_raw()) as *mut Type)
        }
    }

    pub fn of_value<'a, 'mid, 'fid>(value: &'a Value<'cid, 'mid, 'fid>) -> &'a Type<'cid> {
        unsafe {
            &*(LLVMTypeOf(value.as_raw()) as *mut Type)
        }
    }

    pub fn pointer<'ctx, SubType: DerivesFrom<Type<'cid>>>(inner: &'ctx SubType, address_space: c_uint) -> &'ctx PointerType<'cid, SubType> {
        unsafe {
            &*(LLVMPointerType(upcast(inner).as_raw(), address_space) as *mut PointerType<SubType>)
        }
    }

    pub fn array<'ctx, SubType: DerivesFrom<Type<'cid>>>(inner: &'ctx SubType, count: c_uint) -> &'ctx ArrayType<'cid, SubType> {
        unsafe {
            &*(LLVMArrayType(upcast(inner).as_raw(), count) as *mut ArrayType<SubType>)
        }
    }

    pub fn function<'ctx>(params: &[&'ctx Type<'cid>], ret: &'ctx Type<'cid>, var_arg: bool) -> &'ctx FunctionType<'cid> {
        unsafe {
            &*(LLVMFunctionType(ret.as_raw(), params.as_ptr() as *mut LLVMTypeRef, params.len() as u32, var_arg as c_int) as *mut FunctionType)
        }
    }

    pub fn i1<'ctx>(context: &'ctx Context<'cid>) -> &'ctx IntegerType<'cid> {
        unsafe {
            &*(LLVMInt1TypeInContext(context.as_raw()) as *mut IntegerType)
        }
    }

    pub fn i8<'ctx>(context: &'ctx Context<'cid>) -> &'ctx IntegerType<'cid> {
        unsafe {
            &*(LLVMInt8TypeInContext(context.as_raw()) as *mut IntegerType)
        }
    }

    pub fn i16<'ctx>(context: &'ctx Context<'cid>) -> &'ctx IntegerType<'cid> {
        unsafe {
            &*(LLVMInt16TypeInContext(context.as_raw()) as *mut IntegerType)
        }
    }

    pub fn i32<'ctx>(context: &'ctx Context<'cid>) -> &'ctx IntegerType<'cid> {
        unsafe {
            &*(LLVMInt32TypeInContext(context.as_raw()) as *mut IntegerType)
        }
    }

    pub fn i64<'ctx>(context: &'ctx Context<'cid>) -> &'ctx IntegerType<'cid> {
        unsafe {
            &*(LLVMInt64TypeInContext(context.as_raw()) as *mut IntegerType)
        }
    }
}

impl<'cid> FunctionType<'cid> {
    pub fn is_var_arg(&self) -> bool {
        unsafe {
            LLVMIsFunctionVarArg(upcast::<_,Type>(self).as_raw()) != 0
        }
    }
}

impl<'cid, PointeeTy: ?Sized> PointerType<'cid, PointeeTy> {
    pub fn pointee_ty(&self) -> &PointeeTy {
        unsafe {
            &*transmute_copy::<_,*mut PointeeTy>(&LLVMGetElementType(upcast::<_,Type>(self).as_raw()))
        }
    }
}

impl<'cid, ElementTy: ?Sized> ArrayType<'cid, ElementTy> {
    pub fn element_ty(&self) -> &ElementTy {
        unsafe {
            &*transmute_copy::<_,*mut ElementTy>(&LLVMGetElementType(upcast::<_,Type>(self).as_raw()))
        }
    }
}
<commit_msg>Fix missing ?Sized in Type constructors<commit_after>use llvm_sys::prelude::*;
use llvm_sys::core::*;

use libc::{c_int, c_uint};

use std::marker::PhantomData;
use std::mem::transmute_copy;

use id::IdRef;
use opaque::Opaque;
use inheritance::{upcast, DerivesFrom};

use llvm::value::Value;
use llvm::context::Context;

pub struct Type<'cid> {
    _context_id: IdRef<'cid>,
    _opaque: Opaque
}
unsafe impl<'cid> DerivesFrom<Type<'cid>> for Type<'cid> { }

pub struct PointerType<'cid, SubType: ?Sized> {
    _inner: PhantomData<SubType>,
    _super: Type<'cid>
}
unsafe impl<'cid, SubTypeSpecific: DerivesFrom<SubTypeGeneral> + ?Sized, SubTypeGeneral: ?Sized> DerivesFrom<PointerType<'cid, SubTypeGeneral>> for PointerType<'cid, SubTypeSpecific> { }
unsafe impl<'cid, General: ?Sized, SubType: ?Sized> DerivesFrom<General> for PointerType<'cid, SubType> where Type<'cid>: DerivesFrom<General> { }

pub struct ArrayType<'cid, SubType: ?Sized> {
    _inner: PhantomData<SubType>,
    _super: Type<'cid>
}
unsafe impl<'cid, SubTypeSpecific: DerivesFrom<SubTypeGeneral> + ?Sized, SubTypeGeneral: ?Sized> DerivesFrom<ArrayType<'cid, SubTypeGeneral>> for ArrayType<'cid, SubTypeSpecific> { }
unsafe impl<'cid, General: ?Sized, SubType: ?Sized> DerivesFrom<General> for ArrayType<'cid, SubType> where Type<'cid>: DerivesFrom<General> { }

pub struct FunctionType<'cid> {
    _super: Type<'cid>
}
unsafe impl<'cid> DerivesFrom<FunctionType<'cid>> for FunctionType<'cid> { }
unsafe impl<'cid, General: ?Sized> DerivesFrom<General> for FunctionType<'cid> where Type<'cid>: DerivesFrom<General> { }

pub struct IntegerType<'cid> {
    _super: Type<'cid>
}
unsafe impl<'cid> DerivesFrom<IntegerType<'cid>> for IntegerType<'cid> { }
unsafe impl<'cid, General: ?Sized> DerivesFrom<General> for IntegerType<'cid> where Type<'cid>: DerivesFrom<General> { }

impl<'cid> Type<'cid> {
    pub fn dump(&self) {
        unsafe {
            LLVMDumpType(self.as_raw())
        }
    }

    pub fn as_raw(&self) -> LLVMTypeRef {
        self as *const Type as *mut Type as LLVMTypeRef
    }

    pub fn f32<'ctx>(context: &'ctx Context<'cid>) -> &'ctx Type<'cid> {
        unsafe {
            &*(LLVMFloatTypeInContext(context.as_raw()) as *mut Type)
        }
    }

    pub fn f64<'ctx>(context: &'ctx Context<'cid>) -> &'ctx Type<'cid> {
        unsafe {
            &*(LLVMDoubleTypeInContext(context.as_raw()) as *mut Type)
        }
    }

    pub fn void<'ctx>(context: &'ctx Context<'cid>) -> &'ctx Type<'cid> {
        unsafe {
            &*(LLVMVoidTypeInContext(context.as_raw()) as *mut Type)
        }
    }

    pub fn of_value<'a, 'mid, 'fid>(value: &'a Value<'cid, 'mid, 'fid>) -> &'a Type<'cid> {
        unsafe {
            &*(LLVMTypeOf(value.as_raw()) as *mut Type)
        }
    }

    pub fn pointer<'ctx, SubType: DerivesFrom<Type<'cid>> + ?Sized>(inner: &'ctx SubType, address_space: c_uint) -> &'ctx PointerType<'cid, SubType> {
        unsafe {
            &*(LLVMPointerType(upcast(inner).as_raw(), address_space) as *mut PointerType<SubType>)
        }
    }

    pub fn array<'ctx, SubType: DerivesFrom<Type<'cid>> + ?Sized>(inner: &'ctx SubType, count: c_uint) -> &'ctx ArrayType<'cid, SubType> {
        unsafe {
            &*(LLVMArrayType(upcast(inner).as_raw(), count) as *mut ArrayType<SubType>)
        }
    }

    pub fn function<'ctx>(params: &[&'ctx Type<'cid>], ret: &'ctx Type<'cid>, var_arg: bool) -> &'ctx FunctionType<'cid> {
        unsafe {
            &*(LLVMFunctionType(ret.as_raw(), params.as_ptr() as *mut LLVMTypeRef, params.len() as u32, var_arg as c_int) as *mut FunctionType)
        }
    }

    pub fn i1<'ctx>(context: &'ctx Context<'cid>) -> &'ctx IntegerType<'cid> {
        unsafe {
            &*(LLVMInt1TypeInContext(context.as_raw()) as *mut IntegerType)
        }
    }

    pub fn i8<'ctx>(context: &'ctx Context<'cid>) -> &'ctx IntegerType<'cid> {
        unsafe {
            &*(LLVMInt8TypeInContext(context.as_raw()) as *mut IntegerType)
        }
    }

    pub fn i16<'ctx>(context: &'ctx Context<'cid>) -> &'ctx IntegerType<'cid> {
        unsafe {
            &*(LLVMInt16TypeInContext(context.as_raw()) as *mut IntegerType)
        }
    }

    pub fn i32<'ctx>(context: &'ctx Context<'cid>) -> &'ctx IntegerType<'cid> {
        unsafe {
            &*(LLVMInt32TypeInContext(context.as_raw()) as *mut IntegerType)
        }
    }

    pub fn i64<'ctx>(context: &'ctx Context<'cid>) -> &'ctx IntegerType<'cid> {
        unsafe {
            &*(LLVMInt64TypeInContext(context.as_raw()) as *mut IntegerType)
        }
    }
}

impl<'cid> FunctionType<'cid> {
    pub fn is_var_arg(&self) -> bool {
        unsafe {
            LLVMIsFunctionVarArg(upcast::<_,Type>(self).as_raw()) != 0
        }
    }
}

impl<'cid, PointeeTy: ?Sized> PointerType<'cid, PointeeTy> {
    pub fn pointee_ty(&self) -> &PointeeTy {
        unsafe {
            &*transmute_copy::<_,*mut PointeeTy>(&LLVMGetElementType(upcast::<_,Type>(self).as_raw()))
        }
    }
}

impl<'cid, ElementTy: ?Sized> ArrayType<'cid, ElementTy> {
    pub fn element_ty(&self) -> &ElementTy {
        unsafe {
            &*transmute_copy::<_,*mut ElementTy>(&LLVMGetElementType(upcast::<_,Type>(self).as_raw()))
        }
    }
}
<|endoftext|>


<commit_before>//! Functions for creating and modifying messages to send across the message bus.
use std::ops::DerefMut;
use std::cell::RefCell;

use dbus_serialize::types::{Path,Variant,Value,BasicValue,Signature};

use marshal::{Marshal,pad_to_multiple};
use demarshal::{demarshal,DemarshalError};

#[derive(Debug,Default,PartialEq)]
pub struct MessageType(pub u8);
pub const MESSAGE_TYPE_INVALID : MessageType        = MessageType(0);
pub const MESSAGE_TYPE_METHOD_CALL : MessageType    = MessageType(1);
pub const MESSAGE_TYPE_METHOD_RETURN : MessageType  = MessageType(2);
pub const MESSAGE_TYPE_ERROR : MessageType          = MessageType(3);
pub const MESSAGE_TYPE_SIGNAL : MessageType         = MessageType(4);

pub const HEADER_FIELD_INVALID : u8     = 0;
pub const HEADER_FIELD_PATH: u8         = 1;
pub const HEADER_FIELD_INTERFACE: u8    = 2;
pub const HEADER_FIELD_MEMBER: u8       = 3;
pub const HEADER_FIELD_ERROR_NAME: u8   = 4;
pub const HEADER_FIELD_REPLY_SERIAL: u8 = 5;
pub const HEADER_FIELD_DESTINATION: u8  = 6;
pub const HEADER_FIELD_SENDER: u8       = 7;
pub const HEADER_FIELD_SIGNATURE: u8    = 8;

pub const FLAGS_NO_REPLY_EXPECTED : u8  = 1;

#[derive(Debug)]
pub struct HeaderField (
    pub u8,
    pub Variant
);

impl Marshal for HeaderField {
    fn dbus_encode(&self, buf: &mut Vec<u8>) -> usize {
        pad_to_multiple(buf, 8);
        let start_len = buf.len();
        let code = self.0 as u8;
        code.dbus_encode(buf);
        self.1.dbus_encode(buf);
        buf.len() - start_len
    }
    fn get_type(&self) -> String {
        "(yv)".to_owned()
    }
}

/// Represents a received message from the message bus
#[derive(Debug,Default)]
pub struct Message {
    pub big_endian: bool,
    pub message_type: MessageType,
    pub flags: u8,
    pub version: u8,
    pub serial: u32,
    pub headers: Vec<HeaderField>,
    pub body: Vec<u8>,

    body_cache: RefCell<Option<Result<Option<Vec<Value>>, DemarshalError>>>
}

impl Marshal for Message {
    fn dbus_encode (&self, buf: &mut Vec<u8>) -> usize {
        let endian = if self.big_endian { 'B' as u8 } else { 'l' as u8 };
        endian.dbus_encode(buf);
        self.message_type.0.dbus_encode(buf);
        self.flags.dbus_encode(buf);
        self.version.dbus_encode(buf);
        let len : u32 = self.body.len() as u32;
        len.dbus_encode(buf);
        self.serial.dbus_encode(buf);
        self.headers.dbus_encode(buf);
        pad_to_multiple(buf, 8);
        0
    }

    fn get_type (&self) -> String {
        panic!("Don't do that.")
    }
}

/// Create a Message for a D-Bus method call.  Once a Message is created, arguments
/// can be added with Message.add_arg
pub fn create_method_call (dest: &str, path: &str, iface: &str, method: &str) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_METHOD_CALL,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),

        body_cache: RefCell::new(None),
    }.add_header(HEADER_FIELD_DESTINATION,
                 Variant::new(Value::from(dest), "s"))
     .add_header(HEADER_FIELD_PATH,
                 Variant::new(Value::BasicValue(BasicValue::ObjectPath(Path(path.to_owned()))), "o"))
     .add_header(HEADER_FIELD_INTERFACE,
                 Variant::new(Value::from(iface), "s"))
     .add_header(HEADER_FIELD_MEMBER,
                 Variant::new(Value::from(method), "s"))
}

/// Create a Message for a D-Bus method return.  Once created, return values can be added
/// with Message.add_arg
pub fn create_method_return(reply_serial: u32) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_METHOD_RETURN,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),

        body_cache: RefCell::new(None),
    }.add_header(HEADER_FIELD_REPLY_SERIAL,
                 Variant::new(Value::from(reply_serial), "u"))
}

/// Create a Message for a D-Bus error.  Once created, return values can be added
/// with Message.add_arg
pub fn create_error(error_name: &str, reply_serial: u32) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_ERROR,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),

        body_cache: RefCell::new(None),
    }.add_header(HEADER_FIELD_REPLY_SERIAL,
                 Variant::new(Value::from(reply_serial), "u"))
     .add_header(HEADER_FIELD_ERROR_NAME,
                 Variant::new(Value::from(error_name), "s"))
}

/// Create a Message for a D-Bus signal.  Once created, return values can be added
/// with Message.add_arg
pub fn create_signal(path: &str, interface: &str, member: &str) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_SIGNAL,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),

        body_cache: RefCell::new(None),
    }.add_header(HEADER_FIELD_PATH,
                 Variant::new(Value::BasicValue(BasicValue::ObjectPath(Path(path.to_owned()))), "o"))
     .add_header(HEADER_FIELD_INTERFACE,
                 Variant::new(Value::from(interface), "s"))
     .add_header(HEADER_FIELD_MEMBER,
                 Variant::new(Value::from(member), "s"))
}

impl Message {
    /// Add the given argument to the Message.  Accepts anything that implements the Marshal
    /// trait, which is most basic types, as well as the general-purpose
    /// dbus_serialize::types::Value enum.
    ///
    /// Note that these calls can be chained together to add multiple arguments, see the example
    ///
    /// # Examples
    /// ```
    /// dbus_bytestream::message::create_method_call("foo", "/bar", "baz", "bloop")
    ///     .add_arg(&1)
    ///     .add_arg(&"string");
    /// ```
    pub fn add_arg(mut self, arg: &Marshal) -> Message {
        if let None = self.get_header(HEADER_FIELD_SIGNATURE) {
            let value = Value::BasicValue(BasicValue::Signature(Signature("".to_owned())));
            let variant = Variant::new(value, "g");
            self = self.add_header(HEADER_FIELD_SIGNATURE, variant);
        };
        {
            let b : &mut Box<Value> = &mut self.get_header(HEADER_FIELD_SIGNATURE).unwrap().object;
            let val : &mut Value = b.deref_mut();
            match *val {
                Value::BasicValue(BasicValue::Signature(ref mut s)) => s.0.push_str(&arg.get_type()),
                _ => panic!("Garbage in signature field")
            };
        }
        arg.dbus_encode(&mut self.body);
        self
    }

    pub fn get_header(&mut self, name: u8) -> Option<&mut Variant> {
        match self.headers.iter().position(|x| { x.0 == name }) {
            Some(idx) => Some(&mut self.headers[idx].1),
            _ => None
        }
    }

    pub fn add_header(mut self, name: u8, val: Variant) -> Message {
        self.headers.push(HeaderField (name, val));
        self
    }

    /// Get the sequence of Values from out of a Message.  Returns None if the message doesn't have
    /// a body.
    pub fn get_body(&self) -> Result<Option<Vec<Value>>,DemarshalError> {
        if self.body.is_empty() {
            return Ok(None);
        }
        let cached = self.body_cache.borrow().is_some();
        if !cached {
            // Get the signature out of the headers
            let v = match self.headers.iter().position(|x| { x.0 == HEADER_FIELD_SIGNATURE }) {
                Some(idx) => &self.headers[idx].1,
                None => return Ok(None)
            };

            let sigval = match *v.object {
                Value::BasicValue(BasicValue::Signature(ref x)) => x,
                _ => return Ok(None)
            };

            let mut body = self.body.clone();
            let mut sig = "(".to_owned() + &sigval.0 + ")";
            let mut offset = 0;
            *self.body_cache.borrow_mut() = Some((|| {
                match try!(demarshal(&mut body, &mut offset, &mut sig)) {
                    Value::Struct(x) => Ok(Some(x.objects)),
                    x => panic!("Didn't get a struct: {:?}", x)
                }
            })());
        }
        self.body_cache.borrow().as_ref().unwrap().clone()
    }
}

#[test]
fn test_msg () {
    create_method_call("foo", "bar", "baz", "floob")
        .add_arg(&1)
        .add_arg(&2);
}
<commit_msg>message: move mutable get_header call to get_header_mut<commit_after>//! Functions for creating and modifying messages to send across the message bus.
use std::ops::DerefMut;
use std::cell::RefCell;

use dbus_serialize::types::{Path,Variant,Value,BasicValue,Signature};

use marshal::{Marshal,pad_to_multiple};
use demarshal::{demarshal,DemarshalError};

#[derive(Debug,Default,PartialEq)]
pub struct MessageType(pub u8);
pub const MESSAGE_TYPE_INVALID : MessageType        = MessageType(0);
pub const MESSAGE_TYPE_METHOD_CALL : MessageType    = MessageType(1);
pub const MESSAGE_TYPE_METHOD_RETURN : MessageType  = MessageType(2);
pub const MESSAGE_TYPE_ERROR : MessageType          = MessageType(3);
pub const MESSAGE_TYPE_SIGNAL : MessageType         = MessageType(4);

pub const HEADER_FIELD_INVALID : u8     = 0;
pub const HEADER_FIELD_PATH: u8         = 1;
pub const HEADER_FIELD_INTERFACE: u8    = 2;
pub const HEADER_FIELD_MEMBER: u8       = 3;
pub const HEADER_FIELD_ERROR_NAME: u8   = 4;
pub const HEADER_FIELD_REPLY_SERIAL: u8 = 5;
pub const HEADER_FIELD_DESTINATION: u8  = 6;
pub const HEADER_FIELD_SENDER: u8       = 7;
pub const HEADER_FIELD_SIGNATURE: u8    = 8;

pub const FLAGS_NO_REPLY_EXPECTED : u8  = 1;

#[derive(Debug)]
pub struct HeaderField (
    pub u8,
    pub Variant
);

impl Marshal for HeaderField {
    fn dbus_encode(&self, buf: &mut Vec<u8>) -> usize {
        pad_to_multiple(buf, 8);
        let start_len = buf.len();
        let code = self.0 as u8;
        code.dbus_encode(buf);
        self.1.dbus_encode(buf);
        buf.len() - start_len
    }
    fn get_type(&self) -> String {
        "(yv)".to_owned()
    }
}

/// Represents a received message from the message bus
#[derive(Debug,Default)]
pub struct Message {
    pub big_endian: bool,
    pub message_type: MessageType,
    pub flags: u8,
    pub version: u8,
    pub serial: u32,
    pub headers: Vec<HeaderField>,
    pub body: Vec<u8>,

    body_cache: RefCell<Option<Result<Option<Vec<Value>>, DemarshalError>>>
}

impl Marshal for Message {
    fn dbus_encode (&self, buf: &mut Vec<u8>) -> usize {
        let endian = if self.big_endian { 'B' as u8 } else { 'l' as u8 };
        endian.dbus_encode(buf);
        self.message_type.0.dbus_encode(buf);
        self.flags.dbus_encode(buf);
        self.version.dbus_encode(buf);
        let len : u32 = self.body.len() as u32;
        len.dbus_encode(buf);
        self.serial.dbus_encode(buf);
        self.headers.dbus_encode(buf);
        pad_to_multiple(buf, 8);
        0
    }

    fn get_type (&self) -> String {
        panic!("Don't do that.")
    }
}

/// Create a Message for a D-Bus method call.  Once a Message is created, arguments
/// can be added with Message.add_arg
pub fn create_method_call (dest: &str, path: &str, iface: &str, method: &str) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_METHOD_CALL,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),

        body_cache: RefCell::new(None),
    }.add_header(HEADER_FIELD_DESTINATION,
                 Variant::new(Value::from(dest), "s"))
     .add_header(HEADER_FIELD_PATH,
                 Variant::new(Value::BasicValue(BasicValue::ObjectPath(Path(path.to_owned()))), "o"))
     .add_header(HEADER_FIELD_INTERFACE,
                 Variant::new(Value::from(iface), "s"))
     .add_header(HEADER_FIELD_MEMBER,
                 Variant::new(Value::from(method), "s"))
}

/// Create a Message for a D-Bus method return.  Once created, return values can be added
/// with Message.add_arg
pub fn create_method_return(reply_serial: u32) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_METHOD_RETURN,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),

        body_cache: RefCell::new(None),
    }.add_header(HEADER_FIELD_REPLY_SERIAL,
                 Variant::new(Value::from(reply_serial), "u"))
}

/// Create a Message for a D-Bus error.  Once created, return values can be added
/// with Message.add_arg
pub fn create_error(error_name: &str, reply_serial: u32) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_ERROR,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),

        body_cache: RefCell::new(None),
    }.add_header(HEADER_FIELD_REPLY_SERIAL,
                 Variant::new(Value::from(reply_serial), "u"))
     .add_header(HEADER_FIELD_ERROR_NAME,
                 Variant::new(Value::from(error_name), "s"))
}

/// Create a Message for a D-Bus signal.  Once created, return values can be added
/// with Message.add_arg
pub fn create_signal(path: &str, interface: &str, member: &str) -> Message {
    Message {
        big_endian: false,
        message_type: MESSAGE_TYPE_SIGNAL,
        flags: 0,
        version: 1,
        serial: 0,
        headers: Vec::new(),
        body: Vec::new(),

        body_cache: RefCell::new(None),
    }.add_header(HEADER_FIELD_PATH,
                 Variant::new(Value::BasicValue(BasicValue::ObjectPath(Path(path.to_owned()))), "o"))
     .add_header(HEADER_FIELD_INTERFACE,
                 Variant::new(Value::from(interface), "s"))
     .add_header(HEADER_FIELD_MEMBER,
                 Variant::new(Value::from(member), "s"))
}

impl Message {
    /// Add the given argument to the Message.  Accepts anything that implements the Marshal
    /// trait, which is most basic types, as well as the general-purpose
    /// dbus_serialize::types::Value enum.
    ///
    /// Note that these calls can be chained together to add multiple arguments, see the example
    ///
    /// # Examples
    /// ```
    /// dbus_bytestream::message::create_method_call("foo", "/bar", "baz", "bloop")
    ///     .add_arg(&1)
    ///     .add_arg(&"string");
    /// ```
    pub fn add_arg(mut self, arg: &Marshal) -> Message {
        if let None = self.get_header(HEADER_FIELD_SIGNATURE) {
            let value = Value::BasicValue(BasicValue::Signature(Signature("".to_owned())));
            let variant = Variant::new(value, "g");
            self = self.add_header(HEADER_FIELD_SIGNATURE, variant);
        };
        {
            let b : &mut Box<Value> = &mut self.get_header_mut(HEADER_FIELD_SIGNATURE).unwrap().object;
            let val : &mut Value = b.deref_mut();
            match *val {
                Value::BasicValue(BasicValue::Signature(ref mut s)) => s.0.push_str(&arg.get_type()),
                _ => panic!("Garbage in signature field")
            };
        }
        arg.dbus_encode(&mut self.body);
        self
    }

    pub fn get_header(&self, name: u8) -> Option<&Variant> {
        self.headers.iter().position(|x| { x.0 == name })
            .map(|idx| &self.headers[idx].1)
    }

    pub fn get_header_mut(&mut self, name: u8) -> Option<&mut Variant> {
        match self.headers.iter().position(|x| { x.0 == name }) {
            Some(idx) => Some(&mut self.headers[idx].1),
            _ => None
        }
    }

    pub fn add_header(mut self, name: u8, val: Variant) -> Message {
        self.headers.push(HeaderField (name, val));
        self
    }

    /// Get the sequence of Values from out of a Message.  Returns None if the message doesn't have
    /// a body.
    pub fn get_body(&self) -> Result<Option<Vec<Value>>,DemarshalError> {
        if self.body.is_empty() {
            return Ok(None);
        }
        let cached = self.body_cache.borrow().is_some();
        if !cached {
            // Get the signature out of the headers
            let v = match self.headers.iter().position(|x| { x.0 == HEADER_FIELD_SIGNATURE }) {
                Some(idx) => &self.headers[idx].1,
                None => return Ok(None)
            };

            let sigval = match *v.object {
                Value::BasicValue(BasicValue::Signature(ref x)) => x,
                _ => return Ok(None)
            };

            let mut body = self.body.clone();
            let mut sig = "(".to_owned() + &sigval.0 + ")";
            let mut offset = 0;
            *self.body_cache.borrow_mut() = Some((|| {
                match try!(demarshal(&mut body, &mut offset, &mut sig)) {
                    Value::Struct(x) => Ok(Some(x.objects)),
                    x => panic!("Didn't get a struct: {:?}", x)
                }
            })());
        }
        self.body_cache.borrow().as_ref().unwrap().clone()
    }
}

#[test]
fn test_msg () {
    create_method_call("foo", "bar", "baz", "floob")
        .add_arg(&1)
        .add_arg(&2);
}
<|endoftext|>