<commit_before>//! Support for reading WASM files.
//!
//! Provides `WasmFile` and related types which implement the `Object` trait.
//!
//! Currently implements the minimum required to access DWARF debugging information.
#[cfg(feature = "compression")]
use alloc::borrow::Cow;
use alloc::boxed::Box;
use alloc::vec::Vec;
use core::marker::PhantomData;
use core::{fmt, slice, str};
use target_lexicon::Architecture;
use wasmparser as wp;

use crate::read::{
    self, FileFlags, Object, ObjectSection, ObjectSegment, Relocation, SectionFlags, SectionIndex,
    SectionKind, Symbol, SymbolFlags, SymbolIndex, SymbolKind, SymbolMap, SymbolScope,
    SymbolSection,
};

const SECTION_CUSTOM: u8 = 0;
const SECTION_TYPE: u8 = 1;
const SECTION_IMPORT: u8 = 2;
const SECTION_FUNCTION: u8 = 3;
const SECTION_TABLE: u8 = 4;
const SECTION_MEMORY: u8 = 5;
const SECTION_GLOBAL: u8 = 6;
const SECTION_EXPORT: u8 = 7;
const SECTION_START: u8 = 8;
const SECTION_ELEMENT: u8 = 9;
const SECTION_CODE: u8 = 10;
const SECTION_DATA: u8 = 11;
const SECTION_DATA_COUNT: u8 = 12;
// Update this constant when adding new section id:
const MAX_SECTION_ID: u8 = SECTION_DATA_COUNT;

/// A WebAssembly object file.
#[derive(Debug, Default)]
pub struct WasmFile<'data> {
    // All sections, including custom sections.
    sections: Vec<wp::Section<'data>>,
    // Indices into `sections` of sections with a non-zero id.
    id_sections: Box<[Option<usize>; (MAX_SECTION_ID + 1) as usize]>,
    // Payload of custom section called "name".
    names_data: Option<&'data [u8]>,
    // Whether the file has DWARF information.
    has_debug_symbols: bool,
}

impl<'data> WasmFile<'data> {
    /// Parse the raw wasm data.
    pub fn parse(data: &'data [u8]) -> Result<Self, &'static str> {
        let module = wp::ModuleReader::new(data).map_err(|_| "Invalid WASM header")?;

        let mut file = WasmFile::default();

        for section in module {
            let section = section.map_err(|_| "Invalid section header")?;

            match section.code {
                wp::SectionCode::Custom { kind, name } => {
                    if kind == wp::CustomSectionKind::Name {
                        file.names_data = Some(section.range().slice(data));
                    } else if name.starts_with(".debug_") {
                        file.has_debug_symbols = true;
                    }
                }
                code => {
                    let id = section_code_to_id(code);
                    file.id_sections[id as usize] = Some(file.sections.len());
                }
            }

            file.sections.push(section);
        }

        Ok(file)
    }
}

impl<'data> read::private::Sealed for WasmFile<'data> {}

impl<'data, 'file> Object<'data, 'file> for WasmFile<'data>
where
    'data: 'file,
{
    type Segment = WasmSegment<'data, 'file>;
    type SegmentIterator = WasmSegmentIterator<'data, 'file>;
    type Section = WasmSection<'data, 'file>;
    type SectionIterator = WasmSectionIterator<'data, 'file>;
    type SymbolIterator = WasmSymbolIterator<'data, 'file>;

    #[inline]
    fn architecture(&self) -> Architecture {
        Architecture::Wasm32
    }

    #[inline]
    fn is_little_endian(&self) -> bool {
        true
    }

    #[inline]
    fn is_64(&self) -> bool {
        false
    }

    fn segments(&'file self) -> Self::SegmentIterator {
        WasmSegmentIterator { file: self }
    }

    #[inline]
    fn entry(&'file self) -> u64 {
        // TODO: Convert start section to an address.
        0
    }

    fn section_by_name(&'file self, section_name: &str) -> Option<WasmSection<'data, 'file>> {
        self.sections()
            .find(|section| section.name() == Some(section_name))
    }

    fn section_by_index(&'file self, index: SectionIndex) -> Option<WasmSection<'data, 'file>> {
        let id_section = self.id_sections.get(index.0)?;
        let section = self.sections.get((*id_section)?)?;
        Some(WasmSection { section })
    }

    fn sections(&'file self) -> Self::SectionIterator {
        WasmSectionIterator {
            sections: self.sections.iter(),
        }
    }

    #[inline]
    fn symbol_by_index(&self, _index: SymbolIndex) -> Option<Symbol<'data>> {
        // WASM doesn't need or support looking up symbols by index.
        None
    }

    fn symbols(&'file self) -> Self::SymbolIterator {
        WasmSymbolIterator {
            file: self,
            names: self.names_data.and_then(|names| {
                let func = wp::NameSectionReader::new(names, 0)
                    .ok()?
                    .into_iter()
                    .filter_map(|name| name.ok())
                    .find_map(|name| match name {
                        wp::Name::Function(func) => Some(func),
                        _ => None,
                    })?;

                let reader = func.get_map().ok()?;

                Some(NamingIterator { reader }.enumerate())
            }),
        }
    }

    fn dynamic_symbols(&'file self) -> Self::SymbolIterator {
        WasmSymbolIterator {
            file: self,
            names: None,
        }
    }

    fn symbol_map(&self) -> SymbolMap<'data> {
        SymbolMap {
            symbols: Vec::new(),
        }
    }

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

    #[inline]
    fn flags(&self) -> FileFlags {
        FileFlags::None
    }
}

/// An iterator over the segments of a `WasmFile`.
#[derive(Debug)]
pub struct WasmSegmentIterator<'data, 'file> {
    file: &'file WasmFile<'data>,
}

impl<'data, 'file> Iterator for WasmSegmentIterator<'data, 'file> {
    type Item = WasmSegment<'data, 'file>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        None
    }
}

/// A segment of a `WasmFile`.
#[derive(Debug)]
pub struct WasmSegment<'data, 'file> {
    file: &'file WasmFile<'data>,
}

impl<'data, 'file> read::private::Sealed for WasmSegment<'data, 'file> {}

impl<'data, 'file> ObjectSegment<'data> for WasmSegment<'data, 'file> {
    #[inline]
    fn address(&self) -> u64 {
        unreachable!()
    }

    #[inline]
    fn size(&self) -> u64 {
        unreachable!()
    }

    #[inline]
    fn align(&self) -> u64 {
        unreachable!()
    }

    #[inline]
    fn file_range(&self) -> (u64, u64) {
        unreachable!()
    }

    fn data(&self) -> &'data [u8] {
        unreachable!()
    }

    fn data_range(&self, _address: u64, _size: u64) -> Option<&'data [u8]> {
        unreachable!()
    }

    #[inline]
    fn name(&self) -> Option<&str> {
        unreachable!()
    }
}

/// An iterator over the sections of a `WasmFile`.
#[derive(Debug)]
pub struct WasmSectionIterator<'data, 'file> {
    sections: slice::Iter<'file, wp::Section<'data>>,
}

impl<'data, 'file> Iterator for WasmSectionIterator<'data, 'file> {
    type Item = WasmSection<'data, 'file>;

    fn next(&mut self) -> Option<Self::Item> {
        let section = self.sections.next()?;
        Some(WasmSection { section })
    }
}

/// A section of a `WasmFile`.
#[derive(Debug)]
pub struct WasmSection<'data, 'file> {
    section: &'file wp::Section<'data>,
}

impl<'data, 'file> read::private::Sealed for WasmSection<'data, 'file> {}

impl<'data, 'file> ObjectSection<'data> for WasmSection<'data, 'file> {
    type RelocationIterator = WasmRelocationIterator<'data, 'file>;

    #[inline]
    fn index(&self) -> SectionIndex {
        // Note that we treat all custom sections as index 0.
        // This is ok because they are never looked up by index.
        SectionIndex(section_code_to_id(self.section.code) as usize)
    }

    #[inline]
    fn address(&self) -> u64 {
        // TODO: figure out if this should be different for code sections
        0
    }

    #[inline]
    fn size(&self) -> u64 {
        let range = self.section.range();
        (range.end - range.start) as u64
    }

    #[inline]
    fn align(&self) -> u64 {
        1
    }

    #[inline]
    fn file_range(&self) -> Option<(u64, u64)> {
        None
    }

    #[inline]
    fn data(&self) -> &'data [u8] {
        let mut reader = self.section.get_binary_reader();
        // TODO: raise a feature request upstream to be able
        // to get remaining slice from a BinaryReader directly.
        reader.read_bytes(reader.bytes_remaining()).unwrap()
    }

    fn data_range(&self, _address: u64, _size: u64) -> Option<&'data [u8]> {
        unimplemented!()
    }

    #[cfg(feature = "compression")]
    #[inline]
    fn uncompressed_data(&self) -> Option<Cow<'data, [u8]>> {
        Some(Cow::from(self.data()))
    }

    #[inline]
    fn name(&self) -> Option<&str> {
        Some(match self.section.code {
            wp::SectionCode::Custom { name, .. } => name,
            wp::SectionCode::Type => "<type>",
            wp::SectionCode::Import => "<import>",
            wp::SectionCode::Function => "<function>",
            wp::SectionCode::Table => "<table>",
            wp::SectionCode::Memory => "<memory>",
            wp::SectionCode::Global => "<global>",
            wp::SectionCode::Export => "<export>",
            wp::SectionCode::Start => "<start>",
            wp::SectionCode::Element => "<element>",
            wp::SectionCode::Code => "<code>",
            wp::SectionCode::Data => "<data>",
            wp::SectionCode::DataCount => "<data_count>",
        })
    }

    #[inline]
    fn segment_name(&self) -> Option<&str> {
        None
    }

    #[inline]
    fn kind(&self) -> SectionKind {
        SectionKind::Unknown
    }

    #[inline]
    fn relocations(&self) -> WasmRelocationIterator<'data, 'file> {
        WasmRelocationIterator::default()
    }

    #[inline]
    fn flags(&self) -> SectionFlags {
        SectionFlags::None
    }
}

// Upstream NamingReader doesn't have Debug derived,
// so we provide own wrapper that also serves as an Iterator.
struct NamingIterator<'data> {
    reader: wp::NamingReader<'data>,
}

impl fmt::Debug for NamingIterator<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("NamingIterator")
            .field("count", &self.reader.get_count())
            .finish()
    }
}

impl<'data> Iterator for NamingIterator<'data> {
    type Item = wp::Naming<'data>;

    fn next(&mut self) -> Option<Self::Item> {
        self.reader.read().ok()
    }
}

/// An iterator over the symbols of a `WasmFile`.
#[derive(Debug)]
pub struct WasmSymbolIterator<'data, 'file> {
    file: &'file WasmFile<'data>,
    names: Option<core::iter::Enumerate<NamingIterator<'data>>>,
}

impl<'data, 'file> Iterator for WasmSymbolIterator<'data, 'file> {
    type Item = (SymbolIndex, Symbol<'data>);

    fn next(&mut self) -> Option<Self::Item> {
        let (index, naming) = self.names.as_mut()?.next()?;
        Some((
            SymbolIndex(index),
            Symbol {
                name: Some(naming.name),
                address: naming.index as u64,
                size: 0,
                kind: SymbolKind::Text,
                // TODO: maybe treat each function as a section?
                section: SymbolSection::Unknown,
                weak: false,
                scope: SymbolScope::Unknown,
                flags: SymbolFlags::None,
            },
        ))
    }
}

/// An iterator over the relocations in a `WasmSection`.
#[derive(Debug, Default)]
pub struct WasmRelocationIterator<'data, 'file>(PhantomData<(&'data (), &'file ())>);

impl<'data, 'file> Iterator for WasmRelocationIterator<'data, 'file> {
    type Item = (u64, Relocation);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        None
    }
}

fn section_code_to_id(code: wp::SectionCode) -> u8 {
    match code {
        wp::SectionCode::Custom { .. } => SECTION_CUSTOM,
        wp::SectionCode::Type => SECTION_TYPE,
        wp::SectionCode::Import => SECTION_IMPORT,
        wp::SectionCode::Function => SECTION_FUNCTION,
        wp::SectionCode::Table => SECTION_TABLE,
        wp::SectionCode::Memory => SECTION_MEMORY,
        wp::SectionCode::Global => SECTION_GLOBAL,
        wp::SectionCode::Export => SECTION_EXPORT,
        wp::SectionCode::Start => SECTION_START,
        wp::SectionCode::Element => SECTION_ELEMENT,
        wp::SectionCode::Code => SECTION_CODE,
        wp::SectionCode::Data => SECTION_DATA,
        wp::SectionCode::DataCount => SECTION_DATA_COUNT,
    }
}
<commit_msg>Switch from u8 to usize for Wasm section ids<commit_after>//! Support for reading WASM files.
//!
//! Provides `WasmFile` and related types which implement the `Object` trait.
//!
//! Currently implements the minimum required to access DWARF debugging information.
#[cfg(feature = "compression")]
use alloc::borrow::Cow;
use alloc::boxed::Box;
use alloc::vec::Vec;
use core::marker::PhantomData;
use core::{fmt, slice, str};
use target_lexicon::Architecture;
use wasmparser as wp;

use crate::read::{
    self, FileFlags, Object, ObjectSection, ObjectSegment, Relocation, SectionFlags, SectionIndex,
    SectionKind, Symbol, SymbolFlags, SymbolIndex, SymbolKind, SymbolMap, SymbolScope,
    SymbolSection,
};

const SECTION_CUSTOM: usize = 0;
const SECTION_TYPE: usize = 1;
const SECTION_IMPORT: usize = 2;
const SECTION_FUNCTION: usize = 3;
const SECTION_TABLE: usize = 4;
const SECTION_MEMORY: usize = 5;
const SECTION_GLOBAL: usize = 6;
const SECTION_EXPORT: usize = 7;
const SECTION_START: usize = 8;
const SECTION_ELEMENT: usize = 9;
const SECTION_CODE: usize = 10;
const SECTION_DATA: usize = 11;
const SECTION_DATA_COUNT: usize = 12;
// Update this constant when adding new section id:
const MAX_SECTION_ID: usize = SECTION_DATA_COUNT;

/// A WebAssembly object file.
#[derive(Debug, Default)]
pub struct WasmFile<'data> {
    // All sections, including custom sections.
    sections: Vec<wp::Section<'data>>,
    // Indices into `sections` of sections with a non-zero id.
    id_sections: Box<[Option<usize>; MAX_SECTION_ID + 1]>,
    // Payload of custom section called "name".
    names_data: Option<&'data [u8]>,
    // Whether the file has DWARF information.
    has_debug_symbols: bool,
}

impl<'data> WasmFile<'data> {
    /// Parse the raw wasm data.
    pub fn parse(data: &'data [u8]) -> Result<Self, &'static str> {
        let module = wp::ModuleReader::new(data).map_err(|_| "Invalid WASM header")?;

        let mut file = WasmFile::default();

        for section in module {
            let section = section.map_err(|_| "Invalid section header")?;

            match section.code {
                wp::SectionCode::Custom { kind, name } => {
                    if kind == wp::CustomSectionKind::Name {
                        file.names_data = Some(section.range().slice(data));
                    } else if name.starts_with(".debug_") {
                        file.has_debug_symbols = true;
                    }
                }
                code => {
                    let id = section_code_to_id(code);
                    file.id_sections[id] = Some(file.sections.len());
                }
            }

            file.sections.push(section);
        }

        Ok(file)
    }
}

impl<'data> read::private::Sealed for WasmFile<'data> {}

impl<'data, 'file> Object<'data, 'file> for WasmFile<'data>
where
    'data: 'file,
{
    type Segment = WasmSegment<'data, 'file>;
    type SegmentIterator = WasmSegmentIterator<'data, 'file>;
    type Section = WasmSection<'data, 'file>;
    type SectionIterator = WasmSectionIterator<'data, 'file>;
    type SymbolIterator = WasmSymbolIterator<'data, 'file>;

    #[inline]
    fn architecture(&self) -> Architecture {
        Architecture::Wasm32
    }

    #[inline]
    fn is_little_endian(&self) -> bool {
        true
    }

    #[inline]
    fn is_64(&self) -> bool {
        false
    }

    fn segments(&'file self) -> Self::SegmentIterator {
        WasmSegmentIterator { file: self }
    }

    #[inline]
    fn entry(&'file self) -> u64 {
        // TODO: Convert start section to an address.
        0
    }

    fn section_by_name(&'file self, section_name: &str) -> Option<WasmSection<'data, 'file>> {
        self.sections()
            .find(|section| section.name() == Some(section_name))
    }

    fn section_by_index(&'file self, index: SectionIndex) -> Option<WasmSection<'data, 'file>> {
        let id_section = self.id_sections.get(index.0)?;
        let section = self.sections.get((*id_section)?)?;
        Some(WasmSection { section })
    }

    fn sections(&'file self) -> Self::SectionIterator {
        WasmSectionIterator {
            sections: self.sections.iter(),
        }
    }

    #[inline]
    fn symbol_by_index(&self, _index: SymbolIndex) -> Option<Symbol<'data>> {
        // WASM doesn't need or support looking up symbols by index.
        None
    }

    fn symbols(&'file self) -> Self::SymbolIterator {
        WasmSymbolIterator {
            file: self,
            names: self.names_data.and_then(|names| {
                let func = wp::NameSectionReader::new(names, 0)
                    .ok()?
                    .into_iter()
                    .filter_map(|name| name.ok())
                    .find_map(|name| match name {
                        wp::Name::Function(func) => Some(func),
                        _ => None,
                    })?;

                let reader = func.get_map().ok()?;

                Some(NamingIterator { reader }.enumerate())
            }),
        }
    }

    fn dynamic_symbols(&'file self) -> Self::SymbolIterator {
        WasmSymbolIterator {
            file: self,
            names: None,
        }
    }

    fn symbol_map(&self) -> SymbolMap<'data> {
        SymbolMap {
            symbols: Vec::new(),
        }
    }

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

    #[inline]
    fn flags(&self) -> FileFlags {
        FileFlags::None
    }
}

/// An iterator over the segments of a `WasmFile`.
#[derive(Debug)]
pub struct WasmSegmentIterator<'data, 'file> {
    file: &'file WasmFile<'data>,
}

impl<'data, 'file> Iterator for WasmSegmentIterator<'data, 'file> {
    type Item = WasmSegment<'data, 'file>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        None
    }
}

/// A segment of a `WasmFile`.
#[derive(Debug)]
pub struct WasmSegment<'data, 'file> {
    file: &'file WasmFile<'data>,
}

impl<'data, 'file> read::private::Sealed for WasmSegment<'data, 'file> {}

impl<'data, 'file> ObjectSegment<'data> for WasmSegment<'data, 'file> {
    #[inline]
    fn address(&self) -> u64 {
        unreachable!()
    }

    #[inline]
    fn size(&self) -> u64 {
        unreachable!()
    }

    #[inline]
    fn align(&self) -> u64 {
        unreachable!()
    }

    #[inline]
    fn file_range(&self) -> (u64, u64) {
        unreachable!()
    }

    fn data(&self) -> &'data [u8] {
        unreachable!()
    }

    fn data_range(&self, _address: u64, _size: u64) -> Option<&'data [u8]> {
        unreachable!()
    }

    #[inline]
    fn name(&self) -> Option<&str> {
        unreachable!()
    }
}

/// An iterator over the sections of a `WasmFile`.
#[derive(Debug)]
pub struct WasmSectionIterator<'data, 'file> {
    sections: slice::Iter<'file, wp::Section<'data>>,
}

impl<'data, 'file> Iterator for WasmSectionIterator<'data, 'file> {
    type Item = WasmSection<'data, 'file>;

    fn next(&mut self) -> Option<Self::Item> {
        let section = self.sections.next()?;
        Some(WasmSection { section })
    }
}

/// A section of a `WasmFile`.
#[derive(Debug)]
pub struct WasmSection<'data, 'file> {
    section: &'file wp::Section<'data>,
}

impl<'data, 'file> read::private::Sealed for WasmSection<'data, 'file> {}

impl<'data, 'file> ObjectSection<'data> for WasmSection<'data, 'file> {
    type RelocationIterator = WasmRelocationIterator<'data, 'file>;

    #[inline]
    fn index(&self) -> SectionIndex {
        // Note that we treat all custom sections as index 0.
        // This is ok because they are never looked up by index.
        SectionIndex(section_code_to_id(self.section.code))
    }

    #[inline]
    fn address(&self) -> u64 {
        // TODO: figure out if this should be different for code sections
        0
    }

    #[inline]
    fn size(&self) -> u64 {
        let range = self.section.range();
        (range.end - range.start) as u64
    }

    #[inline]
    fn align(&self) -> u64 {
        1
    }

    #[inline]
    fn file_range(&self) -> Option<(u64, u64)> {
        None
    }

    #[inline]
    fn data(&self) -> &'data [u8] {
        let mut reader = self.section.get_binary_reader();
        // TODO: raise a feature request upstream to be able
        // to get remaining slice from a BinaryReader directly.
        reader.read_bytes(reader.bytes_remaining()).unwrap()
    }

    fn data_range(&self, _address: u64, _size: u64) -> Option<&'data [u8]> {
        unimplemented!()
    }

    #[cfg(feature = "compression")]
    #[inline]
    fn uncompressed_data(&self) -> Option<Cow<'data, [u8]>> {
        Some(Cow::from(self.data()))
    }

    #[inline]
    fn name(&self) -> Option<&str> {
        Some(match self.section.code {
            wp::SectionCode::Custom { name, .. } => name,
            wp::SectionCode::Type => "<type>",
            wp::SectionCode::Import => "<import>",
            wp::SectionCode::Function => "<function>",
            wp::SectionCode::Table => "<table>",
            wp::SectionCode::Memory => "<memory>",
            wp::SectionCode::Global => "<global>",
            wp::SectionCode::Export => "<export>",
            wp::SectionCode::Start => "<start>",
            wp::SectionCode::Element => "<element>",
            wp::SectionCode::Code => "<code>",
            wp::SectionCode::Data => "<data>",
            wp::SectionCode::DataCount => "<data_count>",
        })
    }

    #[inline]
    fn segment_name(&self) -> Option<&str> {
        None
    }

    #[inline]
    fn kind(&self) -> SectionKind {
        SectionKind::Unknown
    }

    #[inline]
    fn relocations(&self) -> WasmRelocationIterator<'data, 'file> {
        WasmRelocationIterator::default()
    }

    #[inline]
    fn flags(&self) -> SectionFlags {
        SectionFlags::None
    }
}

// Upstream NamingReader doesn't have Debug derived,
// so we provide own wrapper that also serves as an Iterator.
struct NamingIterator<'data> {
    reader: wp::NamingReader<'data>,
}

impl fmt::Debug for NamingIterator<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("NamingIterator")
            .field("count", &self.reader.get_count())
            .finish()
    }
}

impl<'data> Iterator for NamingIterator<'data> {
    type Item = wp::Naming<'data>;

    fn next(&mut self) -> Option<Self::Item> {
        self.reader.read().ok()
    }
}

/// An iterator over the symbols of a `WasmFile`.
#[derive(Debug)]
pub struct WasmSymbolIterator<'data, 'file> {
    file: &'file WasmFile<'data>,
    names: Option<core::iter::Enumerate<NamingIterator<'data>>>,
}

impl<'data, 'file> Iterator for WasmSymbolIterator<'data, 'file> {
    type Item = (SymbolIndex, Symbol<'data>);

    fn next(&mut self) -> Option<Self::Item> {
        let (index, naming) = self.names.as_mut()?.next()?;
        Some((
            SymbolIndex(index),
            Symbol {
                name: Some(naming.name),
                address: naming.index as u64,
                size: 0,
                kind: SymbolKind::Text,
                // TODO: maybe treat each function as a section?
                section: SymbolSection::Unknown,
                weak: false,
                scope: SymbolScope::Unknown,
                flags: SymbolFlags::None,
            },
        ))
    }
}

/// An iterator over the relocations in a `WasmSection`.
#[derive(Debug, Default)]
pub struct WasmRelocationIterator<'data, 'file>(PhantomData<(&'data (), &'file ())>);

impl<'data, 'file> Iterator for WasmRelocationIterator<'data, 'file> {
    type Item = (u64, Relocation);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        None
    }
}

fn section_code_to_id(code: wp::SectionCode) -> usize {
    match code {
        wp::SectionCode::Custom { .. } => SECTION_CUSTOM,
        wp::SectionCode::Type => SECTION_TYPE,
        wp::SectionCode::Import => SECTION_IMPORT,
        wp::SectionCode::Function => SECTION_FUNCTION,
        wp::SectionCode::Table => SECTION_TABLE,
        wp::SectionCode::Memory => SECTION_MEMORY,
        wp::SectionCode::Global => SECTION_GLOBAL,
        wp::SectionCode::Export => SECTION_EXPORT,
        wp::SectionCode::Start => SECTION_START,
        wp::SectionCode::Element => SECTION_ELEMENT,
        wp::SectionCode::Code => SECTION_CODE,
        wp::SectionCode::Data => SECTION_DATA,
        wp::SectionCode::DataCount => SECTION_DATA_COUNT,
    }
}
<|endoftext|>


<commit_before>use std::sync::{Arc};
use std::net::IpAddr;
use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicIsize, Ordering};

use mioco::sync::Mutex;
use trust_dns::op::{Message, ResponseCode, Query};
use trust_dns::rr::{DNSClass, Name, RecordType, RData, Record};
use itertools::Itertools;

use utils::{Result, CloneExt, MessageExt, Future, AsDisplay};
use query::{query_multiple_handle_futures, query as query_one};
use cache::{Cache, RecordSource};
use nscache::NsCache;
use config::Config;
use resolver::{ErrorKind, RcResolver};

#[derive(Eq, PartialEq, Hash, Clone, Debug)]
struct QueryHash(Name, RecordType, DNSClass);
impl<'a> From<&'a Query> for QueryHash {
    fn from(q: &'a Query) -> Self {
        QueryHash(q.get_name().clone(), q.get_query_type(), q.get_query_class())
    }
}

#[derive(Eq, PartialEq, Hash, Clone, Debug)]
enum QueriedItem {
    NS(IpAddr),
    CNAME(Name),
    NSDomain(Name),
    Query(QueryHash),
}
struct RecursiveResolverState {
    queried_items: Mutex<HashSet<QueriedItem>>,
    query: Query,
    parent: RcResolver,
    query_limit: AtomicIsize,
    skip_cache: bool,
}
impl RecursiveResolverState {
    fn new(q: &Query, parent: RcResolver, skip_cache: bool) -> Self {
        let mut queried_items = HashSet::<QueriedItem>::new();
        queried_items.insert(QueriedItem::Query(q.into()));
        RecursiveResolverState {
            queried_items: Mutex::new(queried_items),
            query: q.clone(),
            parent: parent,
            query_limit: AtomicIsize::new(256),
            skip_cache: skip_cache,
        }
    }
    fn new_inner(&self, q: &Query) -> Result<Self> {
        let mut queried_items = HashSet::<QueriedItem>::new();
        // Exclude NS, since we are querying a different domain
        queried_items.extend(
            self.queried_items.lock().unwrap().iter()
            .filter(|x| if let QueriedItem::NS(_) = **x { false } else { true })
            .map(|x| (*x).clone())
        );
        if !queried_items.insert(QueriedItem::Query(q.into())) {
            return Err(ErrorKind::AlreadyQueried.into());
        }
        Ok(RecursiveResolverState {
            queried_items: Mutex::new(queried_items),
            query: q.clone(),
            parent: self.parent.clone(),
            query_limit: AtomicIsize::new(self.query_limit.load(Ordering::Relaxed)),
            skip_cache: self.skip_cache,
        })
    }
}
#[derive(Clone)]
pub struct RecursiveResolver {
    state: Arc<RecursiveResolverState>,
}
impl RecursiveResolver {
    fn get_cache(&self) -> &Cache {
        &self.state.parent.cache
    }
    fn get_ns_cache(&self) -> &NsCache {
        &self.state.parent.ns_cache
    }
    fn get_config(&self) -> &Config {
        &self.state.parent.config
    }
    fn query(&self) -> Result<Message> {
        let ns = self.get_ns_cache().lookup_recursive(self.state.query.get_name());
        self.query_ns_multiple(ns, None)
    }
    fn query_ns_domain(&self, auth_zone: &Name, ns: &Name) -> Result<Message> {
        debug!("Querying IP of NS {} in zone {}", ns, auth_zone);
        let mut query = Query::new();
        query.name(ns.clone())
        .query_type(RecordType::A)
        .query_class(DNSClass::IN);
        // TODO: IPv6
        let result = try!(self.resolve_next(&query));
        let ips = result.get_answers().iter().filter_map(|x| match *x.get_rdata() {
            RData::A {ref address} => Some((IpAddr::V4(*address), x.get_ttl())),
            _ => None,
        }).collect_vec();
        self.get_ns_cache().update(auth_zone, ips.iter().map(|&(ip, ttl)| (ip, ns.clone(), ttl as u64)));
        self.query_ns_multiple(ips.iter().map(|&(ip, _)| ip), None)
    }
    fn query_ns_domain_future(&self, auth_zone: &Name, ns: Name) -> Option<Future<Result<Message>>> {
        if !self.register_query(QueriedItem::NSDomain(ns.clone())) {
            return None;
        }
        let inst = self.clone();
        let auth_zone_clone = auth_zone.clone();
        Some(Future::from_fn(move || inst.query_ns_domain(&auth_zone_clone, &ns)))
    }
    fn query_ns_domain_futures<T>(&self, auth_zone: &Name, ns_iter: T) -> Vec<Future<Result<Message>>> where T: IntoIterator<Item=Name> {
        ns_iter.into_iter()
        .filter_map(|x| self.query_ns_domain_future(auth_zone, x))
        .collect()
    }
    fn query_ns_multiple<TNs, TExtra>(&self, ns: TNs, extra_futures: TExtra) -> Result<Message>
        where TNs: IntoIterator<Item=IpAddr>,
              TExtra: IntoIterator<Item=Future<Result<Message>>>,
    {
        let mut ns_iter = ns.into_iter().peekable();
        let mut extra_iter = extra_futures.into_iter().peekable();
        if ns_iter.peek().is_none() && extra_iter.peek().is_none() {
            return Err(ErrorKind::NoNameserver.into());
        }
        let mut futures: Vec<_> = ns_iter
        .filter_map(|x| self.query_ns_future(&x))
        .collect();
        futures.extend(extra_iter);
        if futures.is_empty() {
            return Err(ErrorKind::LostRace.into());
        }
        query_multiple_handle_futures(&mut futures)
    }
    fn register_query(&self, item: QueriedItem) -> bool {
        let mut guard = self.state.queried_items.lock().unwrap();
        if guard.contains(&item) {
            trace!("Already queried {:?}", item);
            return false;
        }
        if self.state.query_limit.fetch_sub(1, Ordering::Relaxed) <= 0 {
            debug_assert!(false);
            warn!("Already queried too many servers for {:?}, bug or attack?", self.state.query);
            return false;
        }
        guard.insert(item);
        true
    }
    fn query_ns_future(&self, ns: &IpAddr) -> Option<Future<Result<Message>>> {
        if !self.register_query(QueriedItem::NS(*ns)) {
            return None;
        }
        let inst = self.clone();
        let ns_clone = *ns;
        Some(Future::from_fn(move || inst.query_ns(ns_clone)))
    }
    fn query_ns(self, ns: IpAddr) -> Result<Message> {
        let result = try!(query_one(self.state.query.clone(), ns, *self.get_config().query.timeout));
        if result.get_response_code() != ResponseCode::NoError {
            return Ok(result);
        }
        if !result.get_answers().is_empty() {
            return self.handle_normal_resp(result);
        }
        if result.get_name_servers().iter().any(|x| x.get_rr_type() == RecordType::SOA) {
            // Domain exists, but no record
            return Ok(result)
        }
        self.handle_ns_referral(result)
    }
    #[allow(similar_names)]
    fn find_unresolved_cnames(&self, msg: &Message) -> Option<Vec<Record>> {
        let mut name = self.state.query.get_name();
        let query_type = self.state.query.get_query_type();
        let mut unresolved_cnames = Vec::<&Record>::new();
        loop {
            let cnames = {
                let matched_records = || msg.get_answers().iter()
                .filter(|x| x.get_name() == name);
                let have_real_record = matched_records()
                .any(|x| x.get_rr_type() == query_type);
                if have_real_record {
                    return None;
                }
                matched_records().filter(|x| x.get_rr_type() == RecordType::CNAME)
                .collect_vec()
            };
            if cnames.is_empty() {
                break;
            }
            if cnames.len() > 1 {
                debug_assert!(false);
                warn!("Server returned multiple CNAMEs for a single domain");
                // Try to resolve anyways
            }
            unresolved_cnames.push(cnames[0]);
            name = match *cnames[0].get_rdata() {
                RData::CNAME {ref cname} => cname,
                _ => panic!("Record type doesn't match RData"),
            };
        }
        if unresolved_cnames.is_empty() {
            // FIXME: This means the message has unrelated records in answer section?
            return None;
        }
        Some(unresolved_cnames.iter().map(|x| (*x).clone()).collect())
    }
    #[allow(similar_names)]
    fn handle_normal_resp(&self, msg: Message) -> Result<Message> {
        if self.state.query.get_query_type() == RecordType::CNAME {
            // No need to resolve CNAME chain(?)
            return Ok(msg);
        }
        if let Some(unresolved_cnames) = self.find_unresolved_cnames(&msg) {
            let final_record = &unresolved_cnames[unresolved_cnames.len() - 1];
            let next_name = match *final_record.get_rdata() {
                RData::CNAME {ref cname} => cname.clone(),
                _ => panic!("Record type doesn't match RData"),
            };
            if !self.register_query(QueriedItem::CNAME(next_name.clone())) {
                return Err(ErrorKind::LostRace.into());
            }
            debug!("[{}] CNAME referral: {}", self.state.query.get_name(), next_name);
            let mut next_query = self.state.query.clone();
            next_query.name(next_name);
            match self.resolve_next(&next_query) {
                Err(e) => {
                    warn!("Failed to resolve CNAME {} for {}: {:?}",
                          next_query.get_name(), self.state.query.as_disp(), e);
                    return Ok(msg);
                },
                Ok(ref next_msg) if
                    next_msg.get_response_code() != ResponseCode::NoError
                => {
                    warn!("Failed to resolve CNAME {}: Server returned error: {:?}",
                          next_query.get_name(), next_msg.get_response_code());
                    return Ok(msg);
                }
                Ok(next_msg) => {
                    // Simpler way to clear all answers
                    let mut new_msg = msg.truncate();
                    assert!(new_msg.get_queries().is_empty());
                    assert!(new_msg.get_answers().is_empty());
                    assert!(new_msg.get_name_servers().is_empty());
                    assert!(new_msg.get_additional().is_empty());
                    new_msg.truncated(false);
                    new_msg.add_query(self.state.query.clone());
                    unresolved_cnames.iter().cloned().foreach(|x| {
                        new_msg.add_answer(x);
                    });
                    new_msg.add_all_answers(next_msg.get_answers());
                    return Ok(new_msg);
                },
            };
        }
        Ok(msg)
    }
    fn handle_ns_referral(&self, msg: Message) -> Result<Message> {
        let mut ns_domains : HashMap<_, _> = msg.get_name_servers().iter()
        .filter_map(|x| match *x.get_rdata() {
            RData::NS {ref nsdname} => Some((nsdname.clone(), x.get_name().clone())),
            _ => None,
        })
        .collect();
        if ns_domains.is_empty() {
            debug!("Nameserver returned NOERROR and empty answer: {}, {}",
                   msg.get_queries()[0].as_disp(), msg.as_disp());
            return Ok(msg);
        }
        if ns_domains.values().dedup().count() > 1 {
            // FIXME: Is this legal?
            debug_assert!(false);
            warn!("Nameserver returned NS records for multiple zones");
            return Err(ErrorKind::InsaneNsReferral.into());
        }
        let referred_zone = ns_domains.values().next().unwrap().clone();
        if !referred_zone.zone_of(self.state.query.get_name()) {
            warn!("Nameserver returned NS records for incorrect zone");
            return Err(ErrorKind::InsaneNsReferral.into());
        }
        debug!("[{}] NS referral: {}", self.state.query.get_name(), referred_zone);

        // Get IPs of all NSes.
        let ns_items: Vec<_> = msg.get_additional().iter()
        .filter(|x| ns_domains.contains_key(x.get_name()))
        .filter_map(|x| match *x.get_rdata() {
            RData::A {ref address} => Some((
                IpAddr::V4(*address), x.get_name().clone(), x.get_ttl(),
            )),
            RData::AAAA {ref address} => Some((
                IpAddr::V6(*address), x.get_name().clone(), x.get_ttl(),
            )),
            _ => None,
        })
        .collect();
        self.get_ns_cache().update(
            &referred_zone,
            ns_items.iter().map(|&(ip, ref name, ttl)| {
                ns_domains.remove(name);
                (ip, name.clone(), ttl as u64)
            }),
        );
        let orphan_domain_futures = self.query_ns_domain_futures(
            &referred_zone, ns_domains.keys().cloned(),
        );
        self.query_ns_multiple(
            ns_items.iter().map(|x| x.0),
            orphan_domain_futures,
        )
    }
    fn update_cache(&self, msg: &Message) {
        self.get_cache().update_from_message(msg, RecordSource::Recursive);
    }
    pub fn resolve(q: &Query, parent: RcResolver, skip_cache: bool) -> Result<Message> {
        let resolver = RecursiveResolver {
            state: Arc::new(RecursiveResolverState::new(q, parent, skip_cache)),
        };
        let ret = try!(resolver.query());
        resolver.update_cache(&ret);
        Ok(ret)
    }
    fn resolve_next(&self, q: &Query) -> Result<Message> {
        if !self.state.skip_cache {
            if let Some(msg) = self.get_cache().lookup_with_type(
                q.get_name(), q.get_query_type(), |m| m.clone_resp_for(q),
            ) {
                return Ok(msg);
            }
        }
        let resolver = RecursiveResolver {
            state: Arc::new(try!(self.state.new_inner(q))),
        };
        let ret = try!(resolver.query());
        self.update_cache(&ret);
        Ok(ret)
    }
}
<commit_msg>Avoid wasting too much resources on deeply-nested requests<commit_after>use std::sync::{Arc};
use std::net::IpAddr;
use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicIsize, Ordering};

use mioco::sync::Mutex;
use trust_dns::op::{Message, ResponseCode, Query};
use trust_dns::rr::{DNSClass, Name, RecordType, RData, Record};
use itertools::Itertools;

use utils::{Result, CloneExt, MessageExt, Future, AsDisplay};
use query::{query_multiple_handle_futures, query as query_one};
use cache::{Cache, RecordSource};
use nscache::NsCache;
use config::Config;
use resolver::{ErrorKind, RcResolver};

#[derive(Eq, PartialEq, Hash, Clone, Debug)]
struct QueryHash(Name, RecordType, DNSClass);
impl<'a> From<&'a Query> for QueryHash {
    fn from(q: &'a Query) -> Self {
        QueryHash(q.get_name().clone(), q.get_query_type(), q.get_query_class())
    }
}

#[derive(Eq, PartialEq, Hash, Clone, Debug)]
enum QueriedItem {
    NS(IpAddr),
    CNAME(Name),
    NSDomain(Name),
    Query(QueryHash),
}
struct RecursiveResolverState {
    queried_items: Mutex<HashSet<QueriedItem>>,
    query: Query,
    parent: RcResolver,
    query_limit: AtomicIsize,
    skip_cache: bool,
}
impl RecursiveResolverState {
    fn new(q: &Query, parent: RcResolver, skip_cache: bool) -> Self {
        let mut queried_items = HashSet::<QueriedItem>::new();
        queried_items.insert(QueriedItem::Query(q.into()));
        RecursiveResolverState {
            queried_items: Mutex::new(queried_items),
            query: q.clone(),
            parent: parent,
            query_limit: AtomicIsize::new(256),
            skip_cache: skip_cache,
        }
    }
    fn new_inner(&self, q: &Query) -> Result<Self> {
        let mut queried_items = HashSet::<QueriedItem>::new();
        // Exclude NS, since we are querying a different domain
        queried_items.extend(
            self.queried_items.lock().unwrap().iter()
            .filter(|x| if let QueriedItem::NS(_) = **x { false } else { true })
            .map(|x| (*x).clone())
        );
        if !queried_items.insert(QueriedItem::Query(q.into())) {
            return Err(ErrorKind::AlreadyQueried.into());
        }
        Ok(RecursiveResolverState {
            queried_items: Mutex::new(queried_items),
            query: q.clone(),
            parent: self.parent.clone(),
            query_limit: AtomicIsize::new(self.query_limit.load(Ordering::Relaxed)),
            skip_cache: false, // Avoid wasting too much resources on deeply-nested requests
        })
    }
}
#[derive(Clone)]
pub struct RecursiveResolver {
    state: Arc<RecursiveResolverState>,
}
impl RecursiveResolver {
    fn get_cache(&self) -> &Cache {
        &self.state.parent.cache
    }
    fn get_ns_cache(&self) -> &NsCache {
        &self.state.parent.ns_cache
    }
    fn get_config(&self) -> &Config {
        &self.state.parent.config
    }
    fn query(&self) -> Result<Message> {
        let ns = self.get_ns_cache().lookup_recursive(self.state.query.get_name());
        self.query_ns_multiple(ns, None)
    }
    fn query_ns_domain(&self, auth_zone: &Name, ns: &Name) -> Result<Message> {
        debug!("Querying IP of NS {} in zone {}", ns, auth_zone);
        let mut query = Query::new();
        query.name(ns.clone())
        .query_type(RecordType::A)
        .query_class(DNSClass::IN);
        // TODO: IPv6
        let result = try!(self.resolve_next(&query));
        let ips = result.get_answers().iter().filter_map(|x| match *x.get_rdata() {
            RData::A {ref address} => Some((IpAddr::V4(*address), x.get_ttl())),
            _ => None,
        }).collect_vec();
        self.get_ns_cache().update(auth_zone, ips.iter().map(|&(ip, ttl)| (ip, ns.clone(), ttl as u64)));
        self.query_ns_multiple(ips.iter().map(|&(ip, _)| ip), None)
    }
    fn query_ns_domain_future(&self, auth_zone: &Name, ns: Name) -> Option<Future<Result<Message>>> {
        if !self.register_query(QueriedItem::NSDomain(ns.clone())) {
            return None;
        }
        let inst = self.clone();
        let auth_zone_clone = auth_zone.clone();
        Some(Future::from_fn(move || inst.query_ns_domain(&auth_zone_clone, &ns)))
    }
    fn query_ns_domain_futures<T>(&self, auth_zone: &Name, ns_iter: T) -> Vec<Future<Result<Message>>> where T: IntoIterator<Item=Name> {
        ns_iter.into_iter()
        .filter_map(|x| self.query_ns_domain_future(auth_zone, x))
        .collect()
    }
    fn query_ns_multiple<TNs, TExtra>(&self, ns: TNs, extra_futures: TExtra) -> Result<Message>
        where TNs: IntoIterator<Item=IpAddr>,
              TExtra: IntoIterator<Item=Future<Result<Message>>>,
    {
        let mut ns_iter = ns.into_iter().peekable();
        let mut extra_iter = extra_futures.into_iter().peekable();
        if ns_iter.peek().is_none() && extra_iter.peek().is_none() {
            return Err(ErrorKind::NoNameserver.into());
        }
        let mut futures: Vec<_> = ns_iter
        .filter_map(|x| self.query_ns_future(&x))
        .collect();
        futures.extend(extra_iter);
        if futures.is_empty() {
            return Err(ErrorKind::LostRace.into());
        }
        query_multiple_handle_futures(&mut futures)
    }
    fn register_query(&self, item: QueriedItem) -> bool {
        let mut guard = self.state.queried_items.lock().unwrap();
        if guard.contains(&item) {
            trace!("Already queried {:?}", item);
            return false;
        }
        if self.state.query_limit.fetch_sub(1, Ordering::Relaxed) <= 0 {
            debug_assert!(false);
            warn!("Already queried too many servers for {:?}, bug or attack?", self.state.query);
            return false;
        }
        guard.insert(item);
        true
    }
    fn query_ns_future(&self, ns: &IpAddr) -> Option<Future<Result<Message>>> {
        if !self.register_query(QueriedItem::NS(*ns)) {
            return None;
        }
        let inst = self.clone();
        let ns_clone = *ns;
        Some(Future::from_fn(move || inst.query_ns(ns_clone)))
    }
    fn query_ns(self, ns: IpAddr) -> Result<Message> {
        let result = try!(query_one(self.state.query.clone(), ns, *self.get_config().query.timeout));
        if result.get_response_code() != ResponseCode::NoError {
            return Ok(result);
        }
        if !result.get_answers().is_empty() {
            return self.handle_normal_resp(result);
        }
        if result.get_name_servers().iter().any(|x| x.get_rr_type() == RecordType::SOA) {
            // Domain exists, but no record
            return Ok(result)
        }
        self.handle_ns_referral(result)
    }
    #[allow(similar_names)]
    fn find_unresolved_cnames(&self, msg: &Message) -> Option<Vec<Record>> {
        let mut name = self.state.query.get_name();
        let query_type = self.state.query.get_query_type();
        let mut unresolved_cnames = Vec::<&Record>::new();
        loop {
            let cnames = {
                let matched_records = || msg.get_answers().iter()
                .filter(|x| x.get_name() == name);
                let have_real_record = matched_records()
                .any(|x| x.get_rr_type() == query_type);
                if have_real_record {
                    return None;
                }
                matched_records().filter(|x| x.get_rr_type() == RecordType::CNAME)
                .collect_vec()
            };
            if cnames.is_empty() {
                break;
            }
            if cnames.len() > 1 {
                debug_assert!(false);
                warn!("Server returned multiple CNAMEs for a single domain");
                // Try to resolve anyways
            }
            unresolved_cnames.push(cnames[0]);
            name = match *cnames[0].get_rdata() {
                RData::CNAME {ref cname} => cname,
                _ => panic!("Record type doesn't match RData"),
            };
        }
        if unresolved_cnames.is_empty() {
            // FIXME: This means the message has unrelated records in answer section?
            return None;
        }
        Some(unresolved_cnames.iter().map(|x| (*x).clone()).collect())
    }
    #[allow(similar_names)]
    fn handle_normal_resp(&self, msg: Message) -> Result<Message> {
        if self.state.query.get_query_type() == RecordType::CNAME {
            // No need to resolve CNAME chain(?)
            return Ok(msg);
        }
        if let Some(unresolved_cnames) = self.find_unresolved_cnames(&msg) {
            let final_record = &unresolved_cnames[unresolved_cnames.len() - 1];
            let next_name = match *final_record.get_rdata() {
                RData::CNAME {ref cname} => cname.clone(),
                _ => panic!("Record type doesn't match RData"),
            };
            if !self.register_query(QueriedItem::CNAME(next_name.clone())) {
                return Err(ErrorKind::LostRace.into());
            }
            debug!("[{}] CNAME referral: {}", self.state.query.get_name(), next_name);
            let mut next_query = self.state.query.clone();
            next_query.name(next_name);
            match self.resolve_next(&next_query) {
                Err(e) => {
                    warn!("Failed to resolve CNAME {} for {}: {:?}",
                          next_query.get_name(), self.state.query.as_disp(), e);
                    return Ok(msg);
                },
                Ok(ref next_msg) if
                    next_msg.get_response_code() != ResponseCode::NoError
                => {
                    warn!("Failed to resolve CNAME {}: Server returned error: {:?}",
                          next_query.get_name(), next_msg.get_response_code());
                    return Ok(msg);
                }
                Ok(next_msg) => {
                    // Simpler way to clear all answers
                    let mut new_msg = msg.truncate();
                    assert!(new_msg.get_queries().is_empty());
                    assert!(new_msg.get_answers().is_empty());
                    assert!(new_msg.get_name_servers().is_empty());
                    assert!(new_msg.get_additional().is_empty());
                    new_msg.truncated(false);
                    new_msg.add_query(self.state.query.clone());
                    unresolved_cnames.iter().cloned().foreach(|x| {
                        new_msg.add_answer(x);
                    });
                    new_msg.add_all_answers(next_msg.get_answers());
                    return Ok(new_msg);
                },
            };
        }
        Ok(msg)
    }
    fn handle_ns_referral(&self, msg: Message) -> Result<Message> {
        let mut ns_domains : HashMap<_, _> = msg.get_name_servers().iter()
        .filter_map(|x| match *x.get_rdata() {
            RData::NS {ref nsdname} => Some((nsdname.clone(), x.get_name().clone())),
            _ => None,
        })
        .collect();
        if ns_domains.is_empty() {
            debug!("Nameserver returned NOERROR and empty answer: {}, {}",
                   msg.get_queries()[0].as_disp(), msg.as_disp());
            return Ok(msg);
        }
        if ns_domains.values().dedup().count() > 1 {
            // FIXME: Is this legal?
            debug_assert!(false);
            warn!("Nameserver returned NS records for multiple zones");
            return Err(ErrorKind::InsaneNsReferral.into());
        }
        let referred_zone = ns_domains.values().next().unwrap().clone();
        if !referred_zone.zone_of(self.state.query.get_name()) {
            warn!("Nameserver returned NS records for incorrect zone");
            return Err(ErrorKind::InsaneNsReferral.into());
        }
        debug!("[{}] NS referral: {}", self.state.query.get_name(), referred_zone);

        // Get IPs of all NSes.
        let ns_items: Vec<_> = msg.get_additional().iter()
        .filter(|x| ns_domains.contains_key(x.get_name()))
        .filter_map(|x| match *x.get_rdata() {
            RData::A {ref address} => Some((
                IpAddr::V4(*address), x.get_name().clone(), x.get_ttl(),
            )),
            RData::AAAA {ref address} => Some((
                IpAddr::V6(*address), x.get_name().clone(), x.get_ttl(),
            )),
            _ => None,
        })
        .collect();
        self.get_ns_cache().update(
            &referred_zone,
            ns_items.iter().map(|&(ip, ref name, ttl)| {
                ns_domains.remove(name);
                (ip, name.clone(), ttl as u64)
            }),
        );
        let orphan_domain_futures = self.query_ns_domain_futures(
            &referred_zone, ns_domains.keys().cloned(),
        );
        self.query_ns_multiple(
            ns_items.iter().map(|x| x.0),
            orphan_domain_futures,
        )
    }
    fn update_cache(&self, msg: &Message) {
        self.get_cache().update_from_message(msg, RecordSource::Recursive);
    }
    pub fn resolve(q: &Query, parent: RcResolver, skip_cache: bool) -> Result<Message> {
        let resolver = RecursiveResolver {
            state: Arc::new(RecursiveResolverState::new(q, parent, skip_cache)),
        };
        let ret = try!(resolver.query());
        resolver.update_cache(&ret);
        Ok(ret)
    }
    fn resolve_next(&self, q: &Query) -> Result<Message> {
        if !self.state.skip_cache {
            if let Some(msg) = self.get_cache().lookup_with_type(
                q.get_name(), q.get_query_type(), |m| m.clone_resp_for(q),
            ) {
                return Ok(msg);
            }
        }
        let resolver = RecursiveResolver {
            state: Arc::new(try!(self.state.new_inner(q))),
        };
        let ret = try!(resolver.query());
        self.update_cache(&ret);
        Ok(ret)
    }
}
<|endoftext|>


<commit_before>use libc::{c_char};
use std::ffi::{CStr, CString};
use std::{ptr, mem};
use std::u64;
use std::slice;

#[derive(Debug)]
pub struct FileSrc {
    location: Option<String>,
}

#[repr(C)]
pub enum GstFlowReturn {
    Ok = 0,
    NotLinked = -1,
    Flushing = -2,
    Eos = -3,
    NotNegotiated = -4,
    Error = -5,
}

#[repr(C)]
pub enum GBoolean {
    False = 0,
    True = 1,
}

impl FileSrc {
    fn new() -> FileSrc {
        FileSrc { location: None }
    }
}

impl Drop for FileSrc {
    fn drop(&mut self) {
        println!("drop");
    }
}

#[no_mangle]
pub extern "C" fn filesrc_new() -> *mut FileSrc {
    let instance = Box::new(FileSrc::new());
    return Box::into_raw(instance);
}

#[no_mangle]
pub extern "C" fn filesrc_drop(ptr: *mut FileSrc) {
    unsafe { Box::from_raw(ptr) };
}

#[no_mangle]
pub extern "C" fn filesrc_set_location(ptr: *mut FileSrc, location_ptr: *const c_char) {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    if location_ptr.is_null() {
        filesrc.location = None;
    } else {
        let location = unsafe { CStr::from_ptr(location_ptr) };
        filesrc.location = Some(String::from(location.to_str().unwrap()));
    }
}

#[no_mangle]
pub extern "C" fn filesrc_get_location(ptr: *mut FileSrc) -> *mut c_char {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    match filesrc.location {
        Some(ref location) =>
            CString::new(location.clone().into_bytes()).unwrap().into_raw(),
        None =>
            ptr::null_mut()
    }
}

#[no_mangle]
pub extern "C" fn filesrc_fill(ptr: *mut FileSrc, data_ptr: *mut u8, data_len: usize) -> GstFlowReturn {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    println!("{:?}", filesrc);
    let mut data = unsafe { slice::from_raw_parts_mut(data_ptr, data_len) };

    for i in 0..data.len() - 1 {
        data[i] = 1;
    }

    return GstFlowReturn::Ok;
}

#[no_mangle]
pub extern "C" fn filesrc_get_size(ptr: *mut FileSrc) -> u64 {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    return u64::MAX;
}

#[no_mangle]
pub extern "C" fn filesrc_start(ptr: *mut FileSrc) -> GBoolean {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    match filesrc.location {
        None => GBoolean::False,
        Some(_) => GBoolean::True
    }
}

#[no_mangle]
pub extern "C" fn filesrc_stop(ptr: *mut FileSrc) -> GBoolean {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    return GBoolean::True;
}

#[no_mangle]
pub extern "C" fn filesrc_is_seekable(ptr: *mut FileSrc) -> GBoolean {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    return GBoolean::True;
}

<commit_msg>Move everything into the FileSrc implementation<commit_after>use libc::{c_char};
use std::ffi::{CStr, CString};
use std::ptr;
use std::u64;
use std::slice;

#[repr(C)]
pub enum GstFlowReturn {
    Ok = 0,
    NotLinked = -1,
    Flushing = -2,
    Eos = -3,
    NotNegotiated = -4,
    Error = -5,
}

#[repr(C)]
pub enum GBoolean {
    False = 0,
    True = 1,
}

impl GBoolean {
    fn from_bool(v: bool) -> GBoolean {
        match v {
            true => GBoolean::True,
            false => GBoolean::False,
        }
    }
}

#[derive(Debug)]
pub struct FileSrc {
    location: Option<String>,
}

impl FileSrc {
    fn new() -> FileSrc {
        FileSrc { location: None }
    }

    fn set_location(&mut self, location: &Option<String>) {
        self.location = location.clone();
    }

    fn get_location(&self) -> &Option<String> {
        &self.location
    }

    fn is_seekable(&self) -> bool {
        true
    }

    fn get_size(&self) -> u64 {
        u64::MAX
    }

    fn start(&mut self) -> bool {
        match self.location {
            None => false,
            Some(_) => true
        }
    }

    fn stop(&mut self) -> bool {
        true
    }

    fn fill(&mut self, data: &mut [u8]) -> GstFlowReturn {
        for i in 0..data.len() - 1 {
            data[i] = 1;
        }

        return GstFlowReturn::Ok;
    }
}

impl Drop for FileSrc {
    fn drop(&mut self) {
        println!("drop");
    }
}

#[no_mangle]
pub extern "C" fn filesrc_new() -> *mut FileSrc {
    let instance = Box::new(FileSrc::new());
    return Box::into_raw(instance);
}

#[no_mangle]
pub extern "C" fn filesrc_drop(ptr: *mut FileSrc) {
    unsafe { Box::from_raw(ptr) };
}

#[no_mangle]
pub extern "C" fn filesrc_set_location(ptr: *mut FileSrc, location_ptr: *const c_char) {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    if location_ptr.is_null() {
        filesrc.set_location(&None)
    } else {
        let location = unsafe { CStr::from_ptr(location_ptr) };
        filesrc.set_location(&Some(String::from(location.to_str().unwrap())));
    }
}

#[no_mangle]
pub extern "C" fn filesrc_get_location(ptr: *mut FileSrc) -> *mut c_char {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    match *filesrc.get_location() {
        Some(ref location) =>
            CString::new(location.clone().into_bytes()).unwrap().into_raw(),
        None =>
            ptr::null_mut()
    }
}

#[no_mangle]
pub extern "C" fn filesrc_fill(ptr: *mut FileSrc, data_ptr: *mut u8, data_len: usize) -> GstFlowReturn {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    println!("{:?}", filesrc);
    let mut data = unsafe { slice::from_raw_parts_mut(data_ptr, data_len) };
    return filesrc.fill(data);
}

#[no_mangle]
pub extern "C" fn filesrc_get_size(ptr: *mut FileSrc) -> u64 {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    return filesrc.get_size();
}

#[no_mangle]
pub extern "C" fn filesrc_start(ptr: *mut FileSrc) -> GBoolean {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    GBoolean::from_bool(filesrc.start())
}

#[no_mangle]
pub extern "C" fn filesrc_stop(ptr: *mut FileSrc) -> GBoolean {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    GBoolean::from_bool(filesrc.stop())
}

#[no_mangle]
pub extern "C" fn filesrc_is_seekable(ptr: *mut FileSrc) -> GBoolean {
    let filesrc: &mut FileSrc = unsafe { &mut *ptr };

    GBoolean::from_bool(filesrc.is_seekable())
}

<|endoftext|>


<commit_before>// Bitcoin secp256k1 bindings
// Written in 2014 by
//   Dawid Ciężarkiewicz
//   Andrew Poelstra
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # Secp256k1
//! Rust bindings for Pieter Wuille's secp256k1 library, which is used for
//! fast and accurate manipulation of ECDSA signatures on the secp256k1
//! curve. Such signatures are used extensively by the Bitcoin network
//! and its derivatives.
//!

#![crate_type = "lib"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![crate_name = "bitcoin-secp256k1-rs"]
#![comment = "Bindings and wrapper functions for bitcoin secp256k1 library."]
#![feature(phase)]
#![feature(macro_rules)]
#![feature(globs)]  // for tests only

// Coding conventions
#![deny(non_uppercase_statics)]
#![deny(non_camel_case_types)]
#![deny(non_snake_case)]
#![deny(unused_mut)]
#![warn(missing_doc)]

extern crate "rust-crypto" as crypto;

extern crate libc;
extern crate sync;
extern crate test;

use std::io::IoResult;
use std::rand::{OsRng, Rng, SeedableRng};
use libc::c_int;
use sync::one::{Once, ONCE_INIT};

use crypto::fortuna::Fortuna;

mod macros;
pub mod constants;
pub mod ffi;
pub mod key;

/// A tag used for recovering the public key from a compact signature
pub struct RecoveryId(i32);

/// An ECDSA signature
pub struct Signature(uint, [u8, ..constants::MAX_SIGNATURE_SIZE]);

impl Signature {
    /// Converts the signature to a mutable raw pointer suitable for use
    /// with the FFI functions
    #[inline]
    pub fn as_mut_ptr(&mut self) -> *mut u8 {
        let &Signature(_, ref mut data) = self;
        data.as_mut_slice().as_mut_ptr()
    }

    /// Converts the signature to a byte slice suitable for verification
    #[inline]
    pub fn as_slice<'a>(&'a self) -> &'a [u8] {
        let &Signature(len, ref data) = self;
        data.slice_to(len)
    }

    /// Returns the length of the signature
    #[inline]
    pub fn len(&self) -> uint {
        let &Signature(len, _) = self;
        len
    }
}

/// An ECDSA error
#[deriving(PartialEq, Eq, Clone, Show)]
pub enum Error {
    /// Signature failed verification
    IncorrectSignature,
    /// Bad public key
    InvalidPublicKey,
    /// Bad signature
    InvalidSignature,
    /// Bad secret key
    InvalidSecretKey,
    /// Bad nonce
    InvalidNonce,
    /// Boolean-returning function returned the wrong boolean
    Unknown
}

/// Result type
pub type Result<T> = ::std::prelude::Result<T, Error>;

static mut Secp256k1_init : Once = ONCE_INIT;

/// The secp256k1 engine, used to execute all signature operations
pub struct Secp256k1 {
    rng: Fortuna
}

/// Does one-time initialization of the secp256k1 engine. Can be called
/// multiple times, and is called by the `Secp256k1` constructor. This
/// only needs to be called directly if you are using the library without
/// a `Secp256k1` object, e.g. batch key generation through
/// `key::PublicKey::from_secret_key`.
pub fn init() {
    unsafe {
        Secp256k1_init.doit(|| {
            ffi::secp256k1_start();
        });
    }
}

impl Secp256k1 {
    /// Constructs a new secp256k1 engine.
    pub fn new() -> IoResult<Secp256k1> {
        init();
        let mut osrng = try!(OsRng::new());
        let mut seed = [0, ..2048];
        osrng.fill_bytes(seed.as_mut_slice());
        Ok(Secp256k1 { rng: SeedableRng::from_seed(seed.as_slice()) })
    }

    /// Generates a random keypair. Convenience function for `key::SecretKey::new`
    /// and `key::PublicKey::from_secret_key`; call those functions directly for
    /// batch key generation.
    #[inline]
    pub fn generate_keypair(&mut self, compressed: bool)
                            -> (key::SecretKey, key::PublicKey) {
        let sk = key::SecretKey::new(&mut self.rng);
        (sk, key::PublicKey::from_secret_key(&sk, compressed))
    }

    /// Generates a random nonce. Convenience function for `key::Nonce::new`; call
    /// that function directly for batch nonce generation
    #[inline]
    pub fn generate_nonce(&mut self) -> key::Nonce {
        key::Nonce::new(&mut self.rng)
    }

    /// Constructs a signature for `msg` using the secret key `sk` and nonce `nonce`
    pub fn sign(&self, msg: &[u8], sk: &key::SecretKey, nonce: &key::Nonce)
                -> Result<Signature> {
        let mut sig = [0, ..constants::MAX_SIGNATURE_SIZE];
        let mut len = constants::MAX_SIGNATURE_SIZE as c_int;
        unsafe {
            if ffi::secp256k1_ecdsa_sign(msg.as_ptr(), msg.len() as c_int,
                                         sig.as_mut_slice().as_mut_ptr(), &mut len,
                                         sk.as_ptr(), nonce.as_ptr()) != 1 {
                return Err(InvalidNonce);
            }
            // This assertation is probably too late :)
            assert!(len as uint <= constants::MAX_SIGNATURE_SIZE);
        };
        Ok(Signature(len as uint, sig))
    }

    /// Constructs a compact signature for `msg` using the secret key `sk`
    pub fn sign_compact(&self, msg: &[u8], sk: &key::SecretKey, nonce: &key::Nonce)
                        -> Result<(Signature, RecoveryId)> {
        let mut sig = [0, ..constants::MAX_SIGNATURE_SIZE];
        let mut recid = 0;
        unsafe {
            if ffi::secp256k1_ecdsa_sign_compact(msg.as_ptr(), msg.len() as c_int,
                                                 sig.as_mut_slice().as_mut_ptr(), sk.as_ptr(),
                                                 nonce.as_ptr(), &mut recid) != 1 {
                return Err(InvalidNonce);
            }
        };
        Ok((Signature(constants::MAX_COMPACT_SIGNATURE_SIZE, sig), RecoveryId(recid)))
    }

    /// Determines the public key for which `sig` is a valid signature for
    /// `msg`. Returns through the out-pointer `pubkey`.
    pub fn recover_compact(&self, msg: &[u8], sig: &[u8],
                           compressed: bool, recid: RecoveryId)
                            -> Result<key::PublicKey> {
        let mut pk = key::PublicKey::new(compressed);
        let RecoveryId(recid) = recid;

        unsafe {
            let mut len = 0;
            if ffi::secp256k1_ecdsa_recover_compact(msg.as_ptr(), msg.len() as c_int,
                                                    sig.as_ptr(), pk.as_mut_ptr(), &mut len,
                                                    if compressed {1} else {0},
                                                    recid) != 1 {
                return Err(InvalidSignature);
            }
            assert_eq!(len as uint, pk.len());
        };
        Ok(pk)
    }

    /// Checks that `sig` is a valid ECDSA signature for `msg` using the public
    /// key `pubkey`. Returns `Ok(true)` on success.
    #[inline]
    pub fn verify(msg: &[u8], sig: &[u8], pk: &key::PublicKey) -> Result<()> {
        init();  // This is a static function, so we have to init
        let res = unsafe {
            ffi::secp256k1_ecdsa_verify(msg.as_ptr(), msg.len() as c_int,
                                        sig.as_ptr(), sig.len() as c_int,
                                        pk.as_ptr(), pk.len() as c_int)
        };

        match res {
            1 => Ok(()),
            0 => Err(IncorrectSignature),
            -1 => Err(InvalidPublicKey),
            -2 => Err(InvalidSignature),
            _ => unreachable!()
        }
    }
}

#[cfg(test)]
mod tests {
    use std::rand;
    use std::rand::Rng;

    use test::{Bencher, black_box};

    use key::PublicKey;
    use super::Secp256k1;
    use super::{InvalidPublicKey, IncorrectSignature, InvalidSignature};

    #[test]
    fn invalid_pubkey() {
        let mut msg = Vec::from_elem(32, 0u8);
        let sig = Vec::from_elem(32, 0u8);
        let pk = PublicKey::new(true);

        rand::task_rng().fill_bytes(msg.as_mut_slice());

        assert_eq!(Secp256k1::verify(msg.as_mut_slice(), sig.as_slice(), &pk), Err(InvalidPublicKey));
    }

    #[test]
    fn valid_pubkey_uncompressed() {
        let mut s = Secp256k1::new().unwrap();

        let (_, pk) = s.generate_keypair(false);

        let mut msg = Vec::from_elem(32, 0u8);
        let sig = Vec::from_elem(32, 0u8);

        rand::task_rng().fill_bytes(msg.as_mut_slice());

        assert_eq!(Secp256k1::verify(msg.as_mut_slice(), sig.as_slice(), &pk), Err(InvalidSignature));
    }

    #[test]
    fn valid_pubkey_compressed() {
        let mut s = Secp256k1::new().unwrap();

        let (_, pk) = s.generate_keypair(true);
        let mut msg = Vec::from_elem(32, 0u8);
        let sig = Vec::from_elem(32, 0u8);

        rand::task_rng().fill_bytes(msg.as_mut_slice());

        assert_eq!(Secp256k1::verify(msg.as_mut_slice(), sig.as_slice(), &pk), Err(InvalidSignature));
    }

    #[test]
    fn sign() {
        let mut s = Secp256k1::new().unwrap();

        let mut msg = [0u8, ..32];
        rand::task_rng().fill_bytes(msg);

        let (sk, _) = s.generate_keypair(false);
        let nonce = s.generate_nonce();

        s.sign(msg.as_slice(), &sk, &nonce).unwrap();
    }

    #[test]
    fn sign_and_verify() {
        let mut s = Secp256k1::new().unwrap();

        let mut msg = Vec::from_elem(32, 0u8);
        rand::task_rng().fill_bytes(msg.as_mut_slice());

        let (sk, pk) = s.generate_keypair(false);
        let nonce = s.generate_nonce();

        let sig = s.sign(msg.as_slice(), &sk, &nonce).unwrap();

        assert_eq!(Secp256k1::verify(msg.as_slice(), sig.as_slice(), &pk), Ok(()));
    }

    #[test]
    fn sign_and_verify_fail() {
        let mut s = Secp256k1::new().unwrap();

        let mut msg = Vec::from_elem(32, 0u8);
        rand::task_rng().fill_bytes(msg.as_mut_slice());

        let (sk, pk) = s.generate_keypair(false);
        let nonce = s.generate_nonce();

        let sig = s.sign(msg.as_slice(), &sk, &nonce).unwrap();

        rand::task_rng().fill_bytes(msg.as_mut_slice());
        assert_eq!(Secp256k1::verify(msg.as_slice(), sig.as_slice(), &pk), Err(IncorrectSignature));
    }

    #[test]
    fn sign_compact_with_recovery() {
        let mut s = Secp256k1::new().unwrap();

        let mut msg = [0u8, ..32];
        rand::task_rng().fill_bytes(msg.as_mut_slice());

        let (sk, pk) = s.generate_keypair(false);
        let nonce = s.generate_nonce();

        let (sig, recid) = s.sign_compact(msg.as_slice(), &sk, &nonce).unwrap();

        assert_eq!(s.recover_compact(msg.as_slice(), sig.as_slice(), false, recid), Ok(pk));
    }

    #[bench]
    pub fn generate_compressed(bh: &mut Bencher) {
        let mut s = Secp256k1::new().unwrap();
        bh.iter( || {
          let (sk, pk) = s.generate_keypair(true);
          black_box(sk);
          black_box(pk);
        });
    }

    #[bench]
    pub fn generate_uncompressed(bh: &mut Bencher) {
        let mut s = Secp256k1::new().unwrap();
        bh.iter( || {
          let (sk, pk) = s.generate_keypair(false);
          black_box(sk);
          black_box(pk);
        });
    }
}
<commit_msg>Make `verify` accept a `Signature` rather than a slice<commit_after>// Bitcoin secp256k1 bindings
// Written in 2014 by
//   Dawid Ciężarkiewicz
//   Andrew Poelstra
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # Secp256k1
//! Rust bindings for Pieter Wuille's secp256k1 library, which is used for
//! fast and accurate manipulation of ECDSA signatures on the secp256k1
//! curve. Such signatures are used extensively by the Bitcoin network
//! and its derivatives.
//!

#![crate_type = "lib"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![crate_name = "bitcoin-secp256k1-rs"]
#![comment = "Bindings and wrapper functions for bitcoin secp256k1 library."]
#![feature(phase)]
#![feature(macro_rules)]
#![feature(globs)]  // for tests only

// Coding conventions
#![deny(non_uppercase_statics)]
#![deny(non_camel_case_types)]
#![deny(non_snake_case)]
#![deny(unused_mut)]
#![warn(missing_doc)]

extern crate "rust-crypto" as crypto;

extern crate libc;
extern crate sync;
extern crate test;

use std::intrinsics::copy_nonoverlapping_memory;
use std::io::IoResult;
use std::rand::{OsRng, Rng, SeedableRng};
use libc::c_int;
use sync::one::{Once, ONCE_INIT};

use crypto::fortuna::Fortuna;

mod macros;
pub mod constants;
pub mod ffi;
pub mod key;

/// A tag used for recovering the public key from a compact signature
pub struct RecoveryId(i32);

/// An ECDSA signature
pub struct Signature(uint, [u8, ..constants::MAX_SIGNATURE_SIZE]);

impl Signature {
    /// Converts the signature to a raw pointer suitable for use
    /// with the FFI functions
    #[inline]
    pub fn as_ptr(&self) -> *const u8 {
        let &Signature(_, ref data) = self;
        data.as_slice().as_ptr()
    }

    /// Converts the signature to a mutable raw pointer suitable for use
    /// with the FFI functions
    #[inline]
    pub fn as_mut_ptr(&mut self) -> *mut u8 {
        let &Signature(_, ref mut data) = self;
        data.as_mut_slice().as_mut_ptr()
    }

    /// Converts the signature to a byte slice suitable for verification
    #[inline]
    pub fn as_slice<'a>(&'a self) -> &'a [u8] {
        let &Signature(len, ref data) = self;
        data.slice_to(len)
    }

    /// Returns the length of the signature
    #[inline]
    pub fn len(&self) -> uint {
        let &Signature(len, _) = self;
        len
    }

    /// Converts a byte slice to a signature
    #[inline]
    pub fn from_slice(data: &[u8]) -> Result<Signature> {
        if data.len() <= constants::MAX_SIGNATURE_SIZE {
            let mut ret = [0, ..constants::MAX_SIGNATURE_SIZE];
            unsafe {
                copy_nonoverlapping_memory(ret.as_mut_ptr(),
                                           data.as_ptr(),
                                           data.len());
            }
            Ok(Signature(data.len(), ret))
        } else {
            Err(InvalidSignature)
        }
    }
}

/// An ECDSA error
#[deriving(PartialEq, Eq, Clone, Show)]
pub enum Error {
    /// Signature failed verification
    IncorrectSignature,
    /// Bad public key
    InvalidPublicKey,
    /// Bad signature
    InvalidSignature,
    /// Bad secret key
    InvalidSecretKey,
    /// Bad nonce
    InvalidNonce,
    /// Boolean-returning function returned the wrong boolean
    Unknown
}

/// Result type
pub type Result<T> = ::std::prelude::Result<T, Error>;

static mut Secp256k1_init : Once = ONCE_INIT;

/// The secp256k1 engine, used to execute all signature operations
pub struct Secp256k1 {
    rng: Fortuna
}

/// Does one-time initialization of the secp256k1 engine. Can be called
/// multiple times, and is called by the `Secp256k1` constructor. This
/// only needs to be called directly if you are using the library without
/// a `Secp256k1` object, e.g. batch key generation through
/// `key::PublicKey::from_secret_key`.
pub fn init() {
    unsafe {
        Secp256k1_init.doit(|| {
            ffi::secp256k1_start();
        });
    }
}

impl Secp256k1 {
    /// Constructs a new secp256k1 engine.
    pub fn new() -> IoResult<Secp256k1> {
        init();
        let mut osrng = try!(OsRng::new());
        let mut seed = [0, ..2048];
        osrng.fill_bytes(seed.as_mut_slice());
        Ok(Secp256k1 { rng: SeedableRng::from_seed(seed.as_slice()) })
    }

    /// Generates a random keypair. Convenience function for `key::SecretKey::new`
    /// and `key::PublicKey::from_secret_key`; call those functions directly for
    /// batch key generation.
    #[inline]
    pub fn generate_keypair(&mut self, compressed: bool)
                            -> (key::SecretKey, key::PublicKey) {
        let sk = key::SecretKey::new(&mut self.rng);
        (sk, key::PublicKey::from_secret_key(&sk, compressed))
    }

    /// Generates a random nonce. Convenience function for `key::Nonce::new`; call
    /// that function directly for batch nonce generation
    #[inline]
    pub fn generate_nonce(&mut self) -> key::Nonce {
        key::Nonce::new(&mut self.rng)
    }

    /// Constructs a signature for `msg` using the secret key `sk` and nonce `nonce`
    pub fn sign(&self, msg: &[u8], sk: &key::SecretKey, nonce: &key::Nonce)
                -> Result<Signature> {
        let mut sig = [0, ..constants::MAX_SIGNATURE_SIZE];
        let mut len = constants::MAX_SIGNATURE_SIZE as c_int;
        unsafe {
            if ffi::secp256k1_ecdsa_sign(msg.as_ptr(), msg.len() as c_int,
                                         sig.as_mut_slice().as_mut_ptr(), &mut len,
                                         sk.as_ptr(), nonce.as_ptr()) != 1 {
                return Err(InvalidNonce);
            }
            // This assertation is probably too late :)
            assert!(len as uint <= constants::MAX_SIGNATURE_SIZE);
        };
        Ok(Signature(len as uint, sig))
    }

    /// Constructs a compact signature for `msg` using the secret key `sk`
    pub fn sign_compact(&self, msg: &[u8], sk: &key::SecretKey, nonce: &key::Nonce)
                        -> Result<(Signature, RecoveryId)> {
        let mut sig = [0, ..constants::MAX_SIGNATURE_SIZE];
        let mut recid = 0;
        unsafe {
            if ffi::secp256k1_ecdsa_sign_compact(msg.as_ptr(), msg.len() as c_int,
                                                 sig.as_mut_slice().as_mut_ptr(), sk.as_ptr(),
                                                 nonce.as_ptr(), &mut recid) != 1 {
                return Err(InvalidNonce);
            }
        };
        Ok((Signature(constants::MAX_COMPACT_SIGNATURE_SIZE, sig), RecoveryId(recid)))
    }

    /// Determines the public key for which `sig` is a valid signature for
    /// `msg`. Returns through the out-pointer `pubkey`.
    pub fn recover_compact(&self, msg: &[u8], sig: &[u8],
                           compressed: bool, recid: RecoveryId)
                            -> Result<key::PublicKey> {
        let mut pk = key::PublicKey::new(compressed);
        let RecoveryId(recid) = recid;

        unsafe {
            let mut len = 0;
            if ffi::secp256k1_ecdsa_recover_compact(msg.as_ptr(), msg.len() as c_int,
                                                    sig.as_ptr(), pk.as_mut_ptr(), &mut len,
                                                    if compressed {1} else {0},
                                                    recid) != 1 {
                return Err(InvalidSignature);
            }
            assert_eq!(len as uint, pk.len());
        };
        Ok(pk)
    }

    /// Checks that `sig` is a valid ECDSA signature for `msg` using the public
    /// key `pubkey`. Returns `Ok(true)` on success.
    #[inline]
    pub fn verify(msg: &[u8], sig: &Signature, pk: &key::PublicKey) -> Result<()> {
        init();  // This is a static function, so we have to init
        let res = unsafe {
            ffi::secp256k1_ecdsa_verify(msg.as_ptr(), msg.len() as c_int,
                                        sig.as_ptr(), sig.len() as c_int,
                                        pk.as_ptr(), pk.len() as c_int)
        };

        match res {
            1 => Ok(()),
            0 => Err(IncorrectSignature),
            -1 => Err(InvalidPublicKey),
            -2 => Err(InvalidSignature),
            _ => unreachable!()
        }
    }
}

#[cfg(test)]
mod tests {
    use std::rand;
    use std::rand::Rng;

    use test::{Bencher, black_box};

    use key::PublicKey;
    use super::{Secp256k1, Signature};
    use super::{InvalidPublicKey, IncorrectSignature, InvalidSignature};

    #[test]
    fn invalid_pubkey() {
        let mut msg = Vec::from_elem(32, 0u8);
        let sig = Signature::from_slice([0, ..72]).unwrap();
        let pk = PublicKey::new(true);

        rand::task_rng().fill_bytes(msg.as_mut_slice());

        assert_eq!(Secp256k1::verify(msg.as_mut_slice(), &sig, &pk), Err(InvalidPublicKey));
    }

    #[test]
    fn valid_pubkey_uncompressed() {
        let mut s = Secp256k1::new().unwrap();

        let (_, pk) = s.generate_keypair(false);

        let mut msg = Vec::from_elem(32, 0u8);
        let sig = Signature::from_slice([0, ..72]).unwrap();

        rand::task_rng().fill_bytes(msg.as_mut_slice());

        assert_eq!(Secp256k1::verify(msg.as_mut_slice(), &sig, &pk), Err(InvalidSignature));
    }

    #[test]
    fn valid_pubkey_compressed() {
        let mut s = Secp256k1::new().unwrap();

        let (_, pk) = s.generate_keypair(true);
        let mut msg = Vec::from_elem(32, 0u8);
        let sig = Signature::from_slice([0, ..72]).unwrap();

        rand::task_rng().fill_bytes(msg.as_mut_slice());

        assert_eq!(Secp256k1::verify(msg.as_mut_slice(), &sig, &pk), Err(InvalidSignature));
    }

    #[test]
    fn sign() {
        let mut s = Secp256k1::new().unwrap();

        let mut msg = [0u8, ..32];
        rand::task_rng().fill_bytes(msg);

        let (sk, _) = s.generate_keypair(false);
        let nonce = s.generate_nonce();

        s.sign(msg.as_slice(), &sk, &nonce).unwrap();
    }

    #[test]
    fn sign_and_verify() {
        let mut s = Secp256k1::new().unwrap();

        let mut msg = Vec::from_elem(32, 0u8);
        rand::task_rng().fill_bytes(msg.as_mut_slice());

        let (sk, pk) = s.generate_keypair(false);
        let nonce = s.generate_nonce();

        let sig = s.sign(msg.as_slice(), &sk, &nonce).unwrap();

        assert_eq!(Secp256k1::verify(msg.as_slice(), &sig, &pk), Ok(()));
    }

    #[test]
    fn sign_and_verify_fail() {
        let mut s = Secp256k1::new().unwrap();

        let mut msg = Vec::from_elem(32, 0u8);
        rand::task_rng().fill_bytes(msg.as_mut_slice());

        let (sk, pk) = s.generate_keypair(false);
        let nonce = s.generate_nonce();

        let sig = s.sign(msg.as_slice(), &sk, &nonce).unwrap();

        rand::task_rng().fill_bytes(msg.as_mut_slice());
        assert_eq!(Secp256k1::verify(msg.as_slice(), &sig, &pk), Err(IncorrectSignature));
    }

    #[test]
    fn sign_compact_with_recovery() {
        let mut s = Secp256k1::new().unwrap();

        let mut msg = [0u8, ..32];
        rand::task_rng().fill_bytes(msg.as_mut_slice());

        let (sk, pk) = s.generate_keypair(false);
        let nonce = s.generate_nonce();

        let (sig, recid) = s.sign_compact(msg.as_slice(), &sk, &nonce).unwrap();

        assert_eq!(s.recover_compact(msg.as_slice(), sig.as_slice(), false, recid), Ok(pk));
    }

    #[bench]
    pub fn generate_compressed(bh: &mut Bencher) {
        let mut s = Secp256k1::new().unwrap();
        bh.iter( || {
          let (sk, pk) = s.generate_keypair(true);
          black_box(sk);
          black_box(pk);
        });
    }

    #[bench]
    pub fn generate_uncompressed(bh: &mut Bencher) {
        let mut s = Secp256k1::new().unwrap();
        bh.iter( || {
          let (sk, pk) = s.generate_keypair(false);
          black_box(sk);
          black_box(pk);
        });
    }
}
<|endoftext|>