<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

use std;
use std::ffi::CStr;
use std::io::{Read, Seek, SeekFrom};
use std::mem;

use memory_model::{GuestAddress, GuestMemory};
use sys_util;

#[allow(non_camel_case_types)]
mod elf;

#[derive(Debug, PartialEq)]
pub enum Error {
    BigEndianElfOnLittle,
    CommandLineCopy,
    CommandLineOverflow,
    InvalidElfMagicNumber,
    InvalidEntryAddress,
    InvalidProgramHeaderSize,
    InvalidProgramHeaderOffset,
    InvalidProgramHeaderAddress,
    ReadElfHeader,
    ReadKernelImage,
    ReadProgramHeader,
    SeekKernelStart,
    SeekElfStart,
    SeekProgramHeader,
}
pub type Result<T> = std::result::Result<T, Error>;

/// Loads a kernel from a vmlinux elf image to a slice
///
/// # Arguments
///
/// * `guest_mem` - The guest memory region the kernel is written to.
/// * `kernel_image` - Input vmlinux image.
/// * `start_address` - For x86_64, this is the start of the high memory. Kernel should reside above it.
///
/// Returns the entry address of the kernel.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub fn load_kernel<F>(
    guest_mem: &GuestMemory,
    kernel_image: &mut F,
    start_address: usize,
) -> Result<GuestAddress>
where
    F: Read + Seek,
{
    let mut ehdr: elf::Elf64_Ehdr = Default::default();
    kernel_image
        .seek(SeekFrom::Start(0))
        .map_err(|_| Error::SeekElfStart)?;
    unsafe {
        // read_struct is safe when reading a POD struct.  It can be used and dropped without issue.
        sys_util::read_struct(kernel_image, &mut ehdr).map_err(|_| Error::ReadElfHeader)?;
    }

    // Sanity checks
    if ehdr.e_ident[elf::EI_MAG0 as usize] != elf::ELFMAG0 as u8
        || ehdr.e_ident[elf::EI_MAG1 as usize] != elf::ELFMAG1
        || ehdr.e_ident[elf::EI_MAG2 as usize] != elf::ELFMAG2
        || ehdr.e_ident[elf::EI_MAG3 as usize] != elf::ELFMAG3
    {
        return Err(Error::InvalidElfMagicNumber);
    }
    if ehdr.e_ident[elf::EI_DATA as usize] != elf::ELFDATA2LSB as u8 {
        return Err(Error::BigEndianElfOnLittle);
    }
    if ehdr.e_phentsize as usize != mem::size_of::<elf::Elf64_Phdr>() {
        return Err(Error::InvalidProgramHeaderSize);
    }
    if (ehdr.e_phoff as usize) < mem::size_of::<elf::Elf64_Ehdr>() {
        // If the program header is backwards, bail.
        return Err(Error::InvalidProgramHeaderOffset);
    }
    if (ehdr.e_entry as usize) < start_address {
        return Err(Error::InvalidEntryAddress);
    }

    kernel_image
        .seek(SeekFrom::Start(ehdr.e_phoff))
        .map_err(|_| Error::SeekProgramHeader)?;
    let phdrs: Vec<elf::Elf64_Phdr> = unsafe {
        // Reading the structs is safe for a slice of POD structs.
        sys_util::read_struct_slice(kernel_image, ehdr.e_phnum as usize)
            .map_err(|_| Error::ReadProgramHeader)?
    };

    // Read in each section pointed to by the program headers.
    for phdr in &phdrs {
        if (phdr.p_type & elf::PT_LOAD) == 0 || phdr.p_filesz == 0 {
            continue;
        }

        kernel_image
            .seek(SeekFrom::Start(phdr.p_offset))
            .map_err(|_| Error::SeekKernelStart)?;

        let mem_offset = GuestAddress(phdr.p_paddr as usize);
        if mem_offset.offset() < start_address {
            return Err(Error::InvalidProgramHeaderAddress);
        }

        guest_mem
            .read_to_memory(mem_offset, kernel_image, phdr.p_filesz as usize)
            .map_err(|_| Error::ReadKernelImage)?;
    }

    Ok(GuestAddress(ehdr.e_entry as usize))
}

/// Writes the command line string to the given memory slice.
///
/// # Arguments
///
/// * `guest_mem` - A u8 slice that will be partially overwritten by the command line.
/// * `guest_addr` - The address in `guest_mem` at which to load the command line.
/// * `cmdline` - The kernel command line.
pub fn load_cmdline(
    guest_mem: &GuestMemory,
    guest_addr: GuestAddress,
    cmdline: &CStr,
) -> Result<()> {
    let len = cmdline.to_bytes().len();
    if len == 0 {
        return Ok(());
    }

    let end = guest_addr
        .checked_add(len + 1)
        .ok_or(Error::CommandLineOverflow)?; // Extra for null termination.
    if end > guest_mem.end_addr() {
        return Err(Error::CommandLineOverflow)?;
    }

    guest_mem
        .write_slice_at_addr(cmdline.to_bytes_with_nul(), guest_addr)
        .map_err(|_| Error::CommandLineCopy)?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use memory_model::{GuestAddress, GuestMemory};
    use std::io::Cursor;

    const MEM_SIZE: usize = 0x180000;

    fn create_guest_mem() -> GuestMemory {
        GuestMemory::new(&vec![(GuestAddress(0x0), MEM_SIZE)]).unwrap()
    }

    #[test]
    fn cmdline_overflow() {
        let gm = create_guest_mem();
        let cmdline_address = GuestAddress(MEM_SIZE - 5);
        assert_eq!(
            Err(Error::CommandLineOverflow),
            load_cmdline(
                &gm,
                cmdline_address,
                CStr::from_bytes_with_nul(b"12345\0").unwrap(),
            )
        );
    }

    #[test]
    fn cmdline_write_end() {
        let gm = create_guest_mem();
        let mut cmdline_address = GuestAddress(45);
        assert_eq!(
            Ok(()),
            load_cmdline(
                &gm,
                cmdline_address,
                CStr::from_bytes_with_nul(b"1234\0").unwrap(),
            )
        );
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '1' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '2' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '3' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '4' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '\0' as u8);
    }

    // Elf64 image that prints hello world on x86_64.
    fn make_elf_bin() -> Vec<u8> {
        let mut v = Vec::new();
        v.extend_from_slice(include_bytes!("test_elf.bin"));
        v
    }

    #[test]
    fn load_elf() {
        let gm = create_guest_mem();
        let image = make_elf_bin();
        assert_eq!(
            Ok(GuestAddress(0x100000)),
            load_kernel(&gm, &mut Cursor::new(&image), 0)
        );
    }

    #[test]
    fn bad_magic() {
        let gm = create_guest_mem();
        let mut bad_image = make_elf_bin();
        bad_image[0x1] = 0x33;
        assert_eq!(
            Err(Error::InvalidElfMagicNumber),
            load_kernel(&gm, &mut Cursor::new(&bad_image), 0)
        );
    }

    #[test]
    fn bad_endian() {
        // Only little endian is supported
        let gm = create_guest_mem();
        let mut bad_image = make_elf_bin();
        bad_image[0x5] = 2;
        assert_eq!(
            Err(Error::BigEndianElfOnLittle),
            load_kernel(&gm, &mut Cursor::new(&bad_image), 0)
        );
    }

    #[test]
    fn bad_phoff() {
        // program header has to be past the end of the elf header
        let gm = create_guest_mem();
        let mut bad_image = make_elf_bin();
        bad_image[0x20] = 0x10;
        assert_eq!(
            Err(Error::InvalidProgramHeaderOffset),
            load_kernel(&gm, &mut Cursor::new(&bad_image), 0)
        );
    }
}
<commit_msg>kernel_loader: temporary fix for building on arm<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

use std;
use std::ffi::CStr;
use std::io::{Read, Seek, SeekFrom};
use std::mem;

use memory_model::{GuestAddress, GuestMemory};
use sys_util;

#[allow(non_camel_case_types)]
mod elf;

#[derive(Debug, PartialEq)]
pub enum Error {
    BigEndianElfOnLittle,
    CommandLineCopy,
    CommandLineOverflow,
    InvalidElfMagicNumber,
    InvalidEntryAddress,
    InvalidProgramHeaderSize,
    InvalidProgramHeaderOffset,
    InvalidProgramHeaderAddress,
    ReadElfHeader,
    ReadKernelImage,
    ReadProgramHeader,
    SeekKernelStart,
    SeekElfStart,
    SeekProgramHeader,
}
pub type Result<T> = std::result::Result<T, Error>;

/// Loads a kernel from a vmlinux elf image to a slice
///
/// # Arguments
///
/// * `guest_mem` - The guest memory region the kernel is written to.
/// * `kernel_image` - Input vmlinux image.
/// * `start_address` - For x86_64, this is the start of the high memory. Kernel should reside above it.
///
/// Returns the entry address of the kernel.
pub fn load_kernel<F>(
    guest_mem: &GuestMemory,
    kernel_image: &mut F,
    start_address: usize,
) -> Result<GuestAddress>
where
    F: Read + Seek,
{
    let mut ehdr: elf::Elf64_Ehdr = Default::default();
    kernel_image
        .seek(SeekFrom::Start(0))
        .map_err(|_| Error::SeekElfStart)?;
    unsafe {
        // read_struct is safe when reading a POD struct.  It can be used and dropped without issue.
        sys_util::read_struct(kernel_image, &mut ehdr).map_err(|_| Error::ReadElfHeader)?;
    }

    // Sanity checks
    if ehdr.e_ident[elf::EI_MAG0 as usize] != elf::ELFMAG0 as u8
        || ehdr.e_ident[elf::EI_MAG1 as usize] != elf::ELFMAG1
        || ehdr.e_ident[elf::EI_MAG2 as usize] != elf::ELFMAG2
        || ehdr.e_ident[elf::EI_MAG3 as usize] != elf::ELFMAG3
    {
        return Err(Error::InvalidElfMagicNumber);
    }
    if ehdr.e_ident[elf::EI_DATA as usize] != elf::ELFDATA2LSB as u8 {
        return Err(Error::BigEndianElfOnLittle);
    }
    if ehdr.e_phentsize as usize != mem::size_of::<elf::Elf64_Phdr>() {
        return Err(Error::InvalidProgramHeaderSize);
    }
    if (ehdr.e_phoff as usize) < mem::size_of::<elf::Elf64_Ehdr>() {
        // If the program header is backwards, bail.
        return Err(Error::InvalidProgramHeaderOffset);
    }
    if (ehdr.e_entry as usize) < start_address {
        return Err(Error::InvalidEntryAddress);
    }

    kernel_image
        .seek(SeekFrom::Start(ehdr.e_phoff))
        .map_err(|_| Error::SeekProgramHeader)?;
    let phdrs: Vec<elf::Elf64_Phdr> = unsafe {
        // Reading the structs is safe for a slice of POD structs.
        sys_util::read_struct_slice(kernel_image, ehdr.e_phnum as usize)
            .map_err(|_| Error::ReadProgramHeader)?
    };

    // Read in each section pointed to by the program headers.
    for phdr in &phdrs {
        if (phdr.p_type & elf::PT_LOAD) == 0 || phdr.p_filesz == 0 {
            continue;
        }

        kernel_image
            .seek(SeekFrom::Start(phdr.p_offset))
            .map_err(|_| Error::SeekKernelStart)?;

        let mem_offset = GuestAddress(phdr.p_paddr as usize);
        if mem_offset.offset() < start_address {
            return Err(Error::InvalidProgramHeaderAddress);
        }

        guest_mem
            .read_to_memory(mem_offset, kernel_image, phdr.p_filesz as usize)
            .map_err(|_| Error::ReadKernelImage)?;
    }

    Ok(GuestAddress(ehdr.e_entry as usize))
}

/// Writes the command line string to the given memory slice.
///
/// # Arguments
///
/// * `guest_mem` - A u8 slice that will be partially overwritten by the command line.
/// * `guest_addr` - The address in `guest_mem` at which to load the command line.
/// * `cmdline` - The kernel command line.
pub fn load_cmdline(
    guest_mem: &GuestMemory,
    guest_addr: GuestAddress,
    cmdline: &CStr,
) -> Result<()> {
    let len = cmdline.to_bytes().len();
    if len == 0 {
        return Ok(());
    }

    let end = guest_addr
        .checked_add(len + 1)
        .ok_or(Error::CommandLineOverflow)?; // Extra for null termination.
    if end > guest_mem.end_addr() {
        return Err(Error::CommandLineOverflow)?;
    }

    guest_mem
        .write_slice_at_addr(cmdline.to_bytes_with_nul(), guest_addr)
        .map_err(|_| Error::CommandLineCopy)?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use memory_model::{GuestAddress, GuestMemory};
    use std::io::Cursor;

    const MEM_SIZE: usize = 0x180000;

    fn create_guest_mem() -> GuestMemory {
        GuestMemory::new(&vec![(GuestAddress(0x0), MEM_SIZE)]).unwrap()
    }

    #[test]
    fn cmdline_overflow() {
        let gm = create_guest_mem();
        let cmdline_address = GuestAddress(MEM_SIZE - 5);
        assert_eq!(
            Err(Error::CommandLineOverflow),
            load_cmdline(
                &gm,
                cmdline_address,
                CStr::from_bytes_with_nul(b"12345\0").unwrap(),
            )
        );
    }

    #[test]
    fn cmdline_write_end() {
        let gm = create_guest_mem();
        let mut cmdline_address = GuestAddress(45);
        assert_eq!(
            Ok(()),
            load_cmdline(
                &gm,
                cmdline_address,
                CStr::from_bytes_with_nul(b"1234\0").unwrap(),
            )
        );
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '1' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '2' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '3' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '4' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '\0' as u8);
    }

    // Elf64 image that prints hello world on x86_64.
    fn make_elf_bin() -> Vec<u8> {
        let mut v = Vec::new();
        v.extend_from_slice(include_bytes!("test_elf.bin"));
        v
    }

    #[test]
    fn load_elf() {
        let gm = create_guest_mem();
        let image = make_elf_bin();
        assert_eq!(
            Ok(GuestAddress(0x100000)),
            load_kernel(&gm, &mut Cursor::new(&image), 0)
        );
    }

    #[test]
    fn bad_magic() {
        let gm = create_guest_mem();
        let mut bad_image = make_elf_bin();
        bad_image[0x1] = 0x33;
        assert_eq!(
            Err(Error::InvalidElfMagicNumber),
            load_kernel(&gm, &mut Cursor::new(&bad_image), 0)
        );
    }

    #[test]
    fn bad_endian() {
        // Only little endian is supported
        let gm = create_guest_mem();
        let mut bad_image = make_elf_bin();
        bad_image[0x5] = 2;
        assert_eq!(
            Err(Error::BigEndianElfOnLittle),
            load_kernel(&gm, &mut Cursor::new(&bad_image), 0)
        );
    }

    #[test]
    fn bad_phoff() {
        // program header has to be past the end of the elf header
        let gm = create_guest_mem();
        let mut bad_image = make_elf_bin();
        bad_image[0x20] = 0x10;
        assert_eq!(
            Err(Error::InvalidProgramHeaderOffset),
            load_kernel(&gm, &mut Cursor::new(&bad_image), 0)
        );
    }
}
<|endoftext|>


<commit_before>use proc_macro2::{Delimiter, Group, TokenStream, TokenTree};
use quote::{format_ident, quote, ToTokens};
use std::iter::FromIterator;
use syn::parse::{Nothing, ParseStream};
use syn::{
    braced, bracketed, parenthesized, token, Attribute, Error, Ident, Index, LitInt, LitStr,
    Result, Token,
};

pub struct Attrs<'a> {
    pub display: Option<Display<'a>>,
    pub source: Option<&'a Attribute>,
    pub backtrace: Option<&'a Attribute>,
    pub from: Option<&'a Attribute>,
}

#[derive(Clone)]
pub struct Display<'a> {
    pub original: &'a Attribute,
    pub fmt: LitStr,
    pub args: TokenStream,
    pub was_shorthand: bool,
    pub has_bonus_display: bool,
}

pub fn get(input: &[Attribute]) -> Result<Attrs> {
    let mut attrs = Attrs {
        display: None,
        source: None,
        backtrace: None,
        from: None,
    };

    for attr in input {
        if attr.path.is_ident("error") {
            let display = parse_display(attr)?;
            if attrs.display.is_some() {
                return Err(Error::new_spanned(
                    attr,
                    "only one #[error(...)] attribute is allowed",
                ));
            }
            attrs.display = Some(display);
        } else if attr.path.is_ident("source") {
            require_empty_attribute(attr)?;
            if attrs.source.is_some() {
                return Err(Error::new_spanned(attr, "duplicate #[source] attribute"));
            }
            attrs.source = Some(attr);
        } else if attr.path.is_ident("backtrace") {
            require_empty_attribute(attr)?;
            if attrs.backtrace.is_some() {
                return Err(Error::new_spanned(attr, "duplicate #[backtrace] attribute"));
            }
            attrs.backtrace = Some(attr);
        } else if attr.path.is_ident("from") {
            if !attr.tokens.is_empty() {
                // Assume this is meant for derive_more crate or something.
                continue;
            }
            if attrs.from.is_some() {
                return Err(Error::new_spanned(attr, "duplicate #[from] attribute"));
            }
            attrs.from = Some(attr);
        }
    }

    Ok(attrs)
}

fn parse_display(attr: &Attribute) -> Result<Display> {
    attr.parse_args_with(|input: ParseStream| {
        Ok(Display {
            original: attr,
            fmt: input.parse()?,
            args: parse_token_expr(input, false)?,
            was_shorthand: false,
            has_bonus_display: false,
        })
    })
}

fn parse_token_expr(input: ParseStream, mut last_is_comma: bool) -> Result<TokenStream> {
    let mut tokens = Vec::new();
    while !input.is_empty() {
        if last_is_comma && input.peek(Token![.]) {
            if input.peek2(Ident) {
                input.parse::<Token![.]>()?;
                last_is_comma = false;
                continue;
            }
            if input.peek2(LitInt) {
                input.parse::<Token![.]>()?;
                let int: Index = input.parse()?;
                let ident = format_ident!("_{}", int.index, span = int.span);
                tokens.push(TokenTree::Ident(ident));
                last_is_comma = false;
                continue;
            }
        }
        last_is_comma = input.peek(Token![,]);
        let token: TokenTree = if input.peek(token::Paren) {
            let content;
            let delimiter = parenthesized!(content in input);
            let nested = parse_token_expr(&content, true)?;
            let mut group = Group::new(Delimiter::Parenthesis, nested);
            group.set_span(delimiter.span);
            TokenTree::Group(group)
        } else if input.peek(token::Brace) {
            let content;
            let delimiter = braced!(content in input);
            let nested = parse_token_expr(&content, true)?;
            let mut group = Group::new(Delimiter::Brace, nested);
            group.set_span(delimiter.span);
            TokenTree::Group(group)
        } else if input.peek(token::Bracket) {
            let content;
            let delimiter = bracketed!(content in input);
            let nested = parse_token_expr(&content, true)?;
            let mut group = Group::new(Delimiter::Bracket, nested);
            group.set_span(delimiter.span);
            TokenTree::Group(group)
        } else {
            input.parse()?
        };
        tokens.push(token);
    }
    Ok(TokenStream::from_iter(tokens))
}

fn require_empty_attribute(attr: &Attribute) -> Result<()> {
    syn::parse2::<Nothing>(attr.tokens.clone())?;
    Ok(())
}

impl ToTokens for Display<'_> {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let fmt = &self.fmt;
        let args = &self.args;
        if self.was_shorthand && fmt.value() == "{}" {
            let arg = args.clone().into_iter().nth(1).unwrap();
            tokens.extend(quote! {
                std::fmt::Display::fmt(#arg, __formatter)
            });
        } else if self.was_shorthand && fmt.value() == "{:?}" {
            let arg = args.clone().into_iter().nth(1).unwrap();
            tokens.extend(quote! {
                std::fmt::Debug::fmt(#arg, __formatter)
            });
        } else {
            tokens.extend(quote! {
                write!(__formatter, #fmt #args)
            });
        }
    }
}
<commit_msg>Make function for error attribute parsing<commit_after>use proc_macro2::{Delimiter, Group, TokenStream, TokenTree};
use quote::{format_ident, quote, ToTokens};
use std::iter::FromIterator;
use syn::parse::{Nothing, ParseStream};
use syn::{
    braced, bracketed, parenthesized, token, Attribute, Error, Ident, Index, LitInt, LitStr,
    Result, Token,
};

pub struct Attrs<'a> {
    pub display: Option<Display<'a>>,
    pub source: Option<&'a Attribute>,
    pub backtrace: Option<&'a Attribute>,
    pub from: Option<&'a Attribute>,
}

#[derive(Clone)]
pub struct Display<'a> {
    pub original: &'a Attribute,
    pub fmt: LitStr,
    pub args: TokenStream,
    pub was_shorthand: bool,
    pub has_bonus_display: bool,
}

pub fn get(input: &[Attribute]) -> Result<Attrs> {
    let mut attrs = Attrs {
        display: None,
        source: None,
        backtrace: None,
        from: None,
    };

    for attr in input {
        if attr.path.is_ident("error") {
            parse_error_attribute(&mut attrs, attr)?;
        } else if attr.path.is_ident("source") {
            require_empty_attribute(attr)?;
            if attrs.source.is_some() {
                return Err(Error::new_spanned(attr, "duplicate #[source] attribute"));
            }
            attrs.source = Some(attr);
        } else if attr.path.is_ident("backtrace") {
            require_empty_attribute(attr)?;
            if attrs.backtrace.is_some() {
                return Err(Error::new_spanned(attr, "duplicate #[backtrace] attribute"));
            }
            attrs.backtrace = Some(attr);
        } else if attr.path.is_ident("from") {
            if !attr.tokens.is_empty() {
                // Assume this is meant for derive_more crate or something.
                continue;
            }
            if attrs.from.is_some() {
                return Err(Error::new_spanned(attr, "duplicate #[from] attribute"));
            }
            attrs.from = Some(attr);
        }
    }

    Ok(attrs)
}

fn parse_error_attribute<'a>(attrs: &mut Attrs<'a>, attr: &'a Attribute) -> Result<()> {
    attr.parse_args_with(|input: ParseStream| {
        let display = Display {
            original: attr,
            fmt: input.parse()?,
            args: parse_token_expr(input, false)?,
            was_shorthand: false,
            has_bonus_display: false,
        };
        if attrs.display.is_some() {
            return Err(Error::new_spanned(
                attr,
                "only one #[error(...)] attribute is allowed",
            ));
        }
        attrs.display = Some(display);
        Ok(())
    })
}

fn parse_token_expr(input: ParseStream, mut last_is_comma: bool) -> Result<TokenStream> {
    let mut tokens = Vec::new();
    while !input.is_empty() {
        if last_is_comma && input.peek(Token![.]) {
            if input.peek2(Ident) {
                input.parse::<Token![.]>()?;
                last_is_comma = false;
                continue;
            }
            if input.peek2(LitInt) {
                input.parse::<Token![.]>()?;
                let int: Index = input.parse()?;
                let ident = format_ident!("_{}", int.index, span = int.span);
                tokens.push(TokenTree::Ident(ident));
                last_is_comma = false;
                continue;
            }
        }
        last_is_comma = input.peek(Token![,]);
        let token: TokenTree = if input.peek(token::Paren) {
            let content;
            let delimiter = parenthesized!(content in input);
            let nested = parse_token_expr(&content, true)?;
            let mut group = Group::new(Delimiter::Parenthesis, nested);
            group.set_span(delimiter.span);
            TokenTree::Group(group)
        } else if input.peek(token::Brace) {
            let content;
            let delimiter = braced!(content in input);
            let nested = parse_token_expr(&content, true)?;
            let mut group = Group::new(Delimiter::Brace, nested);
            group.set_span(delimiter.span);
            TokenTree::Group(group)
        } else if input.peek(token::Bracket) {
            let content;
            let delimiter = bracketed!(content in input);
            let nested = parse_token_expr(&content, true)?;
            let mut group = Group::new(Delimiter::Bracket, nested);
            group.set_span(delimiter.span);
            TokenTree::Group(group)
        } else {
            input.parse()?
        };
        tokens.push(token);
    }
    Ok(TokenStream::from_iter(tokens))
}

fn require_empty_attribute(attr: &Attribute) -> Result<()> {
    syn::parse2::<Nothing>(attr.tokens.clone())?;
    Ok(())
}

impl ToTokens for Display<'_> {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let fmt = &self.fmt;
        let args = &self.args;
        if self.was_shorthand && fmt.value() == "{}" {
            let arg = args.clone().into_iter().nth(1).unwrap();
            tokens.extend(quote! {
                std::fmt::Display::fmt(#arg, __formatter)
            });
        } else if self.was_shorthand && fmt.value() == "{:?}" {
            let arg = args.clone().into_iter().nth(1).unwrap();
            tokens.extend(quote! {
                std::fmt::Debug::fmt(#arg, __formatter)
            });
        } else {
            tokens.extend(quote! {
                write!(__formatter, #fmt #args)
            });
        }
    }
}
<|endoftext|>


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

/// External linking is a complex implementation to be able to serve a clean and easy-to-use
/// interface.
///
/// Internally, there are no such things as "external links" (plural). Each Entry in the store can
/// only have _one_ external link.
///
/// This library does the following therefor: It allows you to have several external links with one
/// entry, which are internally one file in the store for each link, linked with "internal
/// linking".
///
/// This helps us greatly with deduplication of URLs.
///

use std::ops::DerefMut;
use std::collections::BTreeMap;

use libimagstore::store::Entry;
use libimagstore::store::FileLockEntry;
use libimagstore::store::Store;
use libimagstore::storeid::StoreId;
use libimagstore::storeid::IntoStoreId;
use libimagutil::debug_result::*;

use error::LinkError as LE;
use error::LinkErrorKind as LEK;
use error::MapErrInto;
use result::Result;
use internal::InternalLinker;
use module_path::ModuleEntryPath;

use toml::Value;
use url::Url;
use crypto::sha1::Sha1;
use crypto::digest::Digest;

/// "Link" Type, just an abstraction over `FileLockEntry` to have some convenience internally.
pub struct Link<'a> {
    link: FileLockEntry<'a>
}

impl<'a> Link<'a> {

    pub fn new(fle: FileLockEntry<'a>) -> Link<'a> {
        Link { link: fle }
    }

    /// Get a link Url object from a `FileLockEntry`, ignore errors.
    fn get_link_uri_from_filelockentry(file: &FileLockEntry<'a>) -> Option<Url> {
        file.get_header()
            .read("imag.content.url")
            .ok()
            .and_then(|opt| match opt {
                Some(Value::String(s)) => {
                    debug!("Found url, parsing: {:?}", s);
                    Url::parse(&s[..]).ok()
                },
                _ => None
            })
    }

    pub fn get_url(&self) -> Result<Option<Url>> {
        let opt = self.link
            .get_header()
            .read("imag.content.url");

        match opt {
            Ok(Some(Value::String(s))) => {
                Url::parse(&s[..])
                     .map(Some)
                     .map_err(|e| LE::new(LEK::EntryHeaderReadError, Some(Box::new(e))))
            },
            Ok(None) => Ok(None),
            _ => Err(LE::new(LEK::EntryHeaderReadError, None))
        }
    }

}

pub trait ExternalLinker : InternalLinker {

    /// Get the external links from the implementor object
    fn get_external_links(&self, store: &Store) -> Result<Vec<Url>>;

    /// Set the external links for the implementor object
    fn set_external_links(&mut self, store: &Store, links: Vec<Url>) -> Result<()>;

    /// Add an external link to the implementor object
    fn add_external_link(&mut self, store: &Store, link: Url) -> Result<()>;

    /// Remove an external link from the implementor object
    fn remove_external_link(&mut self, store: &Store, link: Url) -> Result<()>;

}

pub mod iter {
    use libimagutil::debug_result::*;
    use libimagstore::store::Store;

    use internal::Link;
    use internal::iter::LinkIter;
    use error::LinkErrorKind as LEK;
    use error::MapErrInto;
    use result::Result;

    use url::Url;

    /// Helper for building `OnlyExternalIter` and `NoExternalIter`
    ///
    /// The boolean value defines, how to interpret the `is_external_link_storeid()` return value
    /// (here as "pred"):
    ///
    ///     pred | bool | xor | take?
    ///     ---- | ---- | --- | ----
    ///        0 |    0 |   0 |   1
    ///        0 |    1 |   1 |   0
    ///        1 |    1 |   1 |   0
    ///        1 |    1 |   0 |   1
    ///
    /// If `bool` says "take if return value is false", we take the element if the `pred` returns
    /// false... and so on.
    ///
    /// As we can see, the operator between these two operants is `!(a ^ b)`.
    struct ExternalFilterIter(LinkIter, bool);

    impl Iterator for ExternalFilterIter {
        type Item = Link;

        fn next(&mut self) -> Option<Self::Item> {
            use super::is_external_link_storeid;

            while let Some(elem) = self.0.next() {
                if !(self.1 ^ is_external_link_storeid(&elem)) {
                    return Some(elem);
                }
            }
            None
        }
    }

    pub struct OnlyExternalIter(ExternalFilterIter);

    impl OnlyExternalIter {
        pub fn new(li: LinkIter) -> OnlyExternalIter {
            OnlyExternalIter(ExternalFilterIter(li, true))
        }

        pub fn urls<'a>(self, store: &'a Store) -> UrlIter<'a> {
            UrlIter(self, store)
        }
    }

    impl Iterator for OnlyExternalIter {
        type Item = Link;

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

    pub struct NoExternalIter(ExternalFilterIter);

    impl NoExternalIter {
        pub fn new(li: LinkIter) -> NoExternalIter {
            NoExternalIter(ExternalFilterIter(li, false))
        }
    }

    impl Iterator for NoExternalIter {
        type Item = Link;

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

    pub trait OnlyExternalLinks : Sized {
        fn only_external_links(self) -> OnlyExternalIter ;

        fn no_internal_links(self) -> OnlyExternalIter {
            self.only_external_links()
        }
    }

    impl OnlyExternalLinks for LinkIter {
        fn only_external_links(self) -> OnlyExternalIter {
            OnlyExternalIter::new(self)
        }
    }

    pub trait OnlyInternalLinks : Sized {
        fn only_internal_links(self) -> NoExternalIter;

        fn no_external_links(self) -> NoExternalIter {
            self.only_internal_links()
        }
    }

    impl OnlyInternalLinks for LinkIter {
        fn only_internal_links(self) -> NoExternalIter {
            NoExternalIter::new(self)
        }
    }

    pub struct UrlIter<'a>(OnlyExternalIter, &'a Store);

    impl<'a> Iterator for UrlIter<'a> {
        type Item = Result<Url>;

        fn next(&mut self) -> Option<Self::Item> {
            use super::get_external_link_from_file;

            self.0
                .next()
                .map(|id| {
                    debug!("Retrieving entry for id: '{:?}'", id);
                    self.1
                        .retrieve(id.clone())
                        .map_err_into(LEK::StoreReadError)
                        .map_dbg_err(|_| format!("Retrieving entry for id: '{:?}' failed", id))
                        .and_then(|f| {
                            debug!("Store::retrieve({:?}) succeeded", id);
                            debug!("getting external link from file now");
                            get_external_link_from_file(&f)
                                .map_dbg_err(|e| format!("URL -> Err = {:?}", e))
                        })
                })
        }

    }

}

/// Check whether the StoreId starts with `/link/external/`
pub fn is_external_link_storeid(id: &StoreId) -> bool {
    debug!("Checking whether this is a 'links/external/': '{:?}'", id);
    id.local().starts_with("links/external")
}

fn get_external_link_from_file(entry: &FileLockEntry) -> Result<Url> {
    Link::get_link_uri_from_filelockentry(entry) // TODO: Do not hide error by using this function
        .ok_or(LE::new(LEK::StoreReadError, None))
}

/// Implement `ExternalLinker` for `Entry`, hiding the fact that there is no such thing as an external
/// link in an entry, but internal links to other entries which serve as external links, as one
/// entry in the store can only have one external link.
impl ExternalLinker for Entry {

    /// Get the external links from the implementor object
    fn get_external_links(&self, store: &Store) -> Result<Vec<Url>> {
        // Iterate through all internal links and filter for FileLockEntries which live in
        // /link/external/<SHA> -> load these files and get the external link from their headers,
        // put them into the return vector.
        self.get_internal_links()
            .map(|iter| {
                debug!("Getting external links");
                iter.filter(|l| is_external_link_storeid(l))
                    .map(|id| {
                        debug!("Retrieving entry for id: '{:?}'", id);
                        match store.retrieve(id.clone()) {
                            Ok(f) => {
                                debug!("Store::retrieve({:?}) succeeded", id);
                                debug!("getting external link from file now");
                                get_external_link_from_file(&f)
                                    .map_err(|e| { debug!("URL -> Err = {:?}", e); e })
                            },
                            Err(e) => {
                                debug!("Retrieving entry for id: '{:?}' failed", id);
                                Err(LE::new(LEK::StoreReadError, Some(Box::new(e))))
                            }
                        }
                    })
                    .filter_map(|x| x.ok()) // TODO: Do not ignore error here
                    .collect()
            })
            .map_err(|e| LE::new(LEK::StoreReadError, Some(Box::new(e))))
    }

    /// Set the external links for the implementor object
    fn set_external_links(&mut self, store: &Store, links: Vec<Url>) -> Result<()> {
        // Take all the links, generate a SHA sum out of each one, filter out the already existing
        // store entries and store the other URIs in the header of one FileLockEntry each, in
        // the path /link/external/<SHA of the URL>

        debug!("Iterating {} links = {:?}", links.len(), links);
        for link in links { // for all links
            let hash = {
                let mut s = Sha1::new();
                s.input_str(&link.as_str()[..]);
                s.result_str()
            };
            let file_id = try!(
                ModuleEntryPath::new(format!("external/{}", hash)).into_storeid()
                    .map_err_into(LEK::StoreWriteError)
                    .map_dbg_err(|_| {
                        format!("Failed to build StoreId for this hash '{:?}'", hash)
                    })
                );

            debug!("Link    = '{:?}'", link);
            debug!("Hash    = '{:?}'", hash);
            debug!("StoreId = '{:?}'", file_id);

            // retrieve the file from the store, which implicitely creates the entry if it does not
            // exist
            let mut file = try!(store
                .retrieve(file_id.clone())
                .map_err_into(LEK::StoreWriteError)
                .map_dbg_err(|_| {
                    format!("Failed to create or retrieve an file for this link '{:?}'", link)
                }));

            debug!("Generating header content!");
            {
                let mut hdr = file.deref_mut().get_header_mut();

                let mut table = match hdr.read("imag.content") {
                    Ok(Some(Value::Table(table))) => table,
                    Ok(Some(_)) => {
                        warn!("There is a value at 'imag.content' which is not a table.");
                        warn!("Going to override this value");
                        BTreeMap::new()
                    },
                    Ok(None) => BTreeMap::new(),
                    Err(e)   => return Err(LE::new(LEK::StoreWriteError, Some(Box::new(e)))),
                };

                let v = Value::String(link.into_string());

                debug!("setting URL = '{:?}", v);
                table.insert(String::from("url"), v);

                if let Err(e) = hdr.set("imag.content", Value::Table(table)) {
                    return Err(LE::new(LEK::StoreWriteError, Some(Box::new(e))));
                } else {
                    debug!("Setting URL worked");
                }
            }

            // then add an internal link to the new file or return an error if this fails
            if let Err(e) = self.add_internal_link(file.deref_mut()) {
                debug!("Error adding internal link");
                return Err(LE::new(LEK::StoreWriteError, Some(Box::new(e))));
            }
        }
        debug!("Ready iterating");
        Ok(())
    }

    /// Add an external link to the implementor object
    fn add_external_link(&mut self, store: &Store, link: Url) -> Result<()> {
        // get external links, add this one, save them
        debug!("Getting links");
        self.get_external_links(store)
            .and_then(|mut links| {
                debug!("Adding link = '{:?}' to links = {:?}", link, links);
                links.push(link);
                debug!("Setting {} links = {:?}", links.len(), links);
                self.set_external_links(store, links)
            })
    }

    /// Remove an external link from the implementor object
    fn remove_external_link(&mut self, store: &Store, link: Url) -> Result<()> {
        // get external links, remove this one, save them
        self.get_external_links(store)
            .and_then(|links| {
                debug!("Removing link = '{:?}' from links = {:?}", link, links);
                let links = links.into_iter()
                    .filter(|l| l.as_str() != link.as_str())
                    .collect();
                self.set_external_links(store, links)
            })
    }

}

<commit_msg>Add module documentation<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

/// External linking is a complex implementation to be able to serve a clean and easy-to-use
/// interface.
///
/// Internally, there are no such things as "external links" (plural). Each Entry in the store can
/// only have _one_ external link.
///
/// This library does the following therefor: It allows you to have several external links with one
/// entry, which are internally one file in the store for each link, linked with "internal
/// linking".
///
/// This helps us greatly with deduplication of URLs.
///

use std::ops::DerefMut;
use std::collections::BTreeMap;

use libimagstore::store::Entry;
use libimagstore::store::FileLockEntry;
use libimagstore::store::Store;
use libimagstore::storeid::StoreId;
use libimagstore::storeid::IntoStoreId;
use libimagutil::debug_result::*;

use error::LinkError as LE;
use error::LinkErrorKind as LEK;
use error::MapErrInto;
use result::Result;
use internal::InternalLinker;
use module_path::ModuleEntryPath;

use toml::Value;
use url::Url;
use crypto::sha1::Sha1;
use crypto::digest::Digest;

/// "Link" Type, just an abstraction over `FileLockEntry` to have some convenience internally.
pub struct Link<'a> {
    link: FileLockEntry<'a>
}

impl<'a> Link<'a> {

    pub fn new(fle: FileLockEntry<'a>) -> Link<'a> {
        Link { link: fle }
    }

    /// Get a link Url object from a `FileLockEntry`, ignore errors.
    fn get_link_uri_from_filelockentry(file: &FileLockEntry<'a>) -> Option<Url> {
        file.get_header()
            .read("imag.content.url")
            .ok()
            .and_then(|opt| match opt {
                Some(Value::String(s)) => {
                    debug!("Found url, parsing: {:?}", s);
                    Url::parse(&s[..]).ok()
                },
                _ => None
            })
    }

    pub fn get_url(&self) -> Result<Option<Url>> {
        let opt = self.link
            .get_header()
            .read("imag.content.url");

        match opt {
            Ok(Some(Value::String(s))) => {
                Url::parse(&s[..])
                     .map(Some)
                     .map_err(|e| LE::new(LEK::EntryHeaderReadError, Some(Box::new(e))))
            },
            Ok(None) => Ok(None),
            _ => Err(LE::new(LEK::EntryHeaderReadError, None))
        }
    }

}

pub trait ExternalLinker : InternalLinker {

    /// Get the external links from the implementor object
    fn get_external_links(&self, store: &Store) -> Result<Vec<Url>>;

    /// Set the external links for the implementor object
    fn set_external_links(&mut self, store: &Store, links: Vec<Url>) -> Result<()>;

    /// Add an external link to the implementor object
    fn add_external_link(&mut self, store: &Store, link: Url) -> Result<()>;

    /// Remove an external link from the implementor object
    fn remove_external_link(&mut self, store: &Store, link: Url) -> Result<()>;

}

pub mod iter {
    //! Iterator helpers for external linking stuff
    //!
    //! Contains also helpers to filter iterators for external/internal links
    //!

    use libimagutil::debug_result::*;
    use libimagstore::store::Store;

    use internal::Link;
    use internal::iter::LinkIter;
    use error::LinkErrorKind as LEK;
    use error::MapErrInto;
    use result::Result;

    use url::Url;

    /// Helper for building `OnlyExternalIter` and `NoExternalIter`
    ///
    /// The boolean value defines, how to interpret the `is_external_link_storeid()` return value
    /// (here as "pred"):
    ///
    ///     pred | bool | xor | take?
    ///     ---- | ---- | --- | ----
    ///        0 |    0 |   0 |   1
    ///        0 |    1 |   1 |   0
    ///        1 |    1 |   1 |   0
    ///        1 |    1 |   0 |   1
    ///
    /// If `bool` says "take if return value is false", we take the element if the `pred` returns
    /// false... and so on.
    ///
    /// As we can see, the operator between these two operants is `!(a ^ b)`.
    struct ExternalFilterIter(LinkIter, bool);

    impl Iterator for ExternalFilterIter {
        type Item = Link;

        fn next(&mut self) -> Option<Self::Item> {
            use super::is_external_link_storeid;

            while let Some(elem) = self.0.next() {
                if !(self.1 ^ is_external_link_storeid(&elem)) {
                    return Some(elem);
                }
            }
            None
        }
    }

    pub struct OnlyExternalIter(ExternalFilterIter);

    impl OnlyExternalIter {
        pub fn new(li: LinkIter) -> OnlyExternalIter {
            OnlyExternalIter(ExternalFilterIter(li, true))
        }

        pub fn urls<'a>(self, store: &'a Store) -> UrlIter<'a> {
            UrlIter(self, store)
        }
    }

    impl Iterator for OnlyExternalIter {
        type Item = Link;

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

    pub struct NoExternalIter(ExternalFilterIter);

    impl NoExternalIter {
        pub fn new(li: LinkIter) -> NoExternalIter {
            NoExternalIter(ExternalFilterIter(li, false))
        }
    }

    impl Iterator for NoExternalIter {
        type Item = Link;

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

    pub trait OnlyExternalLinks : Sized {
        fn only_external_links(self) -> OnlyExternalIter ;

        fn no_internal_links(self) -> OnlyExternalIter {
            self.only_external_links()
        }
    }

    impl OnlyExternalLinks for LinkIter {
        fn only_external_links(self) -> OnlyExternalIter {
            OnlyExternalIter::new(self)
        }
    }

    pub trait OnlyInternalLinks : Sized {
        fn only_internal_links(self) -> NoExternalIter;

        fn no_external_links(self) -> NoExternalIter {
            self.only_internal_links()
        }
    }

    impl OnlyInternalLinks for LinkIter {
        fn only_internal_links(self) -> NoExternalIter {
            NoExternalIter::new(self)
        }
    }

    pub struct UrlIter<'a>(OnlyExternalIter, &'a Store);

    impl<'a> Iterator for UrlIter<'a> {
        type Item = Result<Url>;

        fn next(&mut self) -> Option<Self::Item> {
            use super::get_external_link_from_file;

            self.0
                .next()
                .map(|id| {
                    debug!("Retrieving entry for id: '{:?}'", id);
                    self.1
                        .retrieve(id.clone())
                        .map_err_into(LEK::StoreReadError)
                        .map_dbg_err(|_| format!("Retrieving entry for id: '{:?}' failed", id))
                        .and_then(|f| {
                            debug!("Store::retrieve({:?}) succeeded", id);
                            debug!("getting external link from file now");
                            get_external_link_from_file(&f)
                                .map_dbg_err(|e| format!("URL -> Err = {:?}", e))
                        })
                })
        }

    }

}

/// Check whether the StoreId starts with `/link/external/`
pub fn is_external_link_storeid(id: &StoreId) -> bool {
    debug!("Checking whether this is a 'links/external/': '{:?}'", id);
    id.local().starts_with("links/external")
}

fn get_external_link_from_file(entry: &FileLockEntry) -> Result<Url> {
    Link::get_link_uri_from_filelockentry(entry) // TODO: Do not hide error by using this function
        .ok_or(LE::new(LEK::StoreReadError, None))
}

/// Implement `ExternalLinker` for `Entry`, hiding the fact that there is no such thing as an external
/// link in an entry, but internal links to other entries which serve as external links, as one
/// entry in the store can only have one external link.
impl ExternalLinker for Entry {

    /// Get the external links from the implementor object
    fn get_external_links(&self, store: &Store) -> Result<Vec<Url>> {
        // Iterate through all internal links and filter for FileLockEntries which live in
        // /link/external/<SHA> -> load these files and get the external link from their headers,
        // put them into the return vector.
        self.get_internal_links()
            .map(|iter| {
                debug!("Getting external links");
                iter.filter(|l| is_external_link_storeid(l))
                    .map(|id| {
                        debug!("Retrieving entry for id: '{:?}'", id);
                        match store.retrieve(id.clone()) {
                            Ok(f) => {
                                debug!("Store::retrieve({:?}) succeeded", id);
                                debug!("getting external link from file now");
                                get_external_link_from_file(&f)
                                    .map_err(|e| { debug!("URL -> Err = {:?}", e); e })
                            },
                            Err(e) => {
                                debug!("Retrieving entry for id: '{:?}' failed", id);
                                Err(LE::new(LEK::StoreReadError, Some(Box::new(e))))
                            }
                        }
                    })
                    .filter_map(|x| x.ok()) // TODO: Do not ignore error here
                    .collect()
            })
            .map_err(|e| LE::new(LEK::StoreReadError, Some(Box::new(e))))
    }

    /// Set the external links for the implementor object
    fn set_external_links(&mut self, store: &Store, links: Vec<Url>) -> Result<()> {
        // Take all the links, generate a SHA sum out of each one, filter out the already existing
        // store entries and store the other URIs in the header of one FileLockEntry each, in
        // the path /link/external/<SHA of the URL>

        debug!("Iterating {} links = {:?}", links.len(), links);
        for link in links { // for all links
            let hash = {
                let mut s = Sha1::new();
                s.input_str(&link.as_str()[..]);
                s.result_str()
            };
            let file_id = try!(
                ModuleEntryPath::new(format!("external/{}", hash)).into_storeid()
                    .map_err_into(LEK::StoreWriteError)
                    .map_dbg_err(|_| {
                        format!("Failed to build StoreId for this hash '{:?}'", hash)
                    })
                );

            debug!("Link    = '{:?}'", link);
            debug!("Hash    = '{:?}'", hash);
            debug!("StoreId = '{:?}'", file_id);

            // retrieve the file from the store, which implicitely creates the entry if it does not
            // exist
            let mut file = try!(store
                .retrieve(file_id.clone())
                .map_err_into(LEK::StoreWriteError)
                .map_dbg_err(|_| {
                    format!("Failed to create or retrieve an file for this link '{:?}'", link)
                }));

            debug!("Generating header content!");
            {
                let mut hdr = file.deref_mut().get_header_mut();

                let mut table = match hdr.read("imag.content") {
                    Ok(Some(Value::Table(table))) => table,
                    Ok(Some(_)) => {
                        warn!("There is a value at 'imag.content' which is not a table.");
                        warn!("Going to override this value");
                        BTreeMap::new()
                    },
                    Ok(None) => BTreeMap::new(),
                    Err(e)   => return Err(LE::new(LEK::StoreWriteError, Some(Box::new(e)))),
                };

                let v = Value::String(link.into_string());

                debug!("setting URL = '{:?}", v);
                table.insert(String::from("url"), v);

                if let Err(e) = hdr.set("imag.content", Value::Table(table)) {
                    return Err(LE::new(LEK::StoreWriteError, Some(Box::new(e))));
                } else {
                    debug!("Setting URL worked");
                }
            }

            // then add an internal link to the new file or return an error if this fails
            if let Err(e) = self.add_internal_link(file.deref_mut()) {
                debug!("Error adding internal link");
                return Err(LE::new(LEK::StoreWriteError, Some(Box::new(e))));
            }
        }
        debug!("Ready iterating");
        Ok(())
    }

    /// Add an external link to the implementor object
    fn add_external_link(&mut self, store: &Store, link: Url) -> Result<()> {
        // get external links, add this one, save them
        debug!("Getting links");
        self.get_external_links(store)
            .and_then(|mut links| {
                debug!("Adding link = '{:?}' to links = {:?}", link, links);
                links.push(link);
                debug!("Setting {} links = {:?}", links.len(), links);
                self.set_external_links(store, links)
            })
    }

    /// Remove an external link from the implementor object
    fn remove_external_link(&mut self, store: &Store, link: Url) -> Result<()> {
        // get external links, remove this one, save them
        self.get_external_links(store)
            .and_then(|links| {
                debug!("Removing link = '{:?}' from links = {:?}", link, links);
                let links = links.into_iter()
                    .filter(|l| l.as_str() != link.as_str())
                    .collect();
                self.set_external_links(store, links)
            })
    }

}

<|endoftext|>


<commit_before>use toml::Value;

use fs2::FileExt;

use libimagstore::hook::Hook;
use libimagstore::hook::accessor::HookDataAccessor as HDA;
use libimagstore::hook::accessor::HookDataAccessorProvider;
use libimagstore::hook::accessor::StoreIdAccessor;
use libimagstore::hook::accessor::MutableHookDataAccessor;
use libimagstore::hook::accessor::NonMutableHookDataAccessor;
use libimagstore::hook::result::HookResult;
use libimagstore::hook::error::{HookError, HookErrorKind, MapErrInto};
use libimagstore::storeid::StoreId;
use libimagstore::store::FileLockEntry;
use libimagstore::store::Entry;

trait EntryFlock {
    fn lock(&self) -> HookResult<()>;
    fn unlock(&self) -> HookResult<()>;
}

impl EntryFlock for Entry {

    fn lock(&self) -> HookResult<()> {
        use std::fs::File;

        self.get_location()
            .clone()
            .into_pathbuf()
            .map_err_into(HookErrorKind::HookExecutionError)
            .and_then(|loc| File::open(loc).map_err_into(HookErrorKind::HookExecutionError))
            .and_then(|file| file.lock_exclusive().map_err_into(HookErrorKind::HookExecutionError))
    }

    fn unlock(&self) -> HookResult<()> {
        use std::fs::File;

        self.get_location()
            .clone()
            .into_pathbuf()
            .map_err_into(HookErrorKind::HookExecutionError)
            .and_then(|loc| File::open(loc).map_err_into(HookErrorKind::HookExecutionError))
            .and_then(|file| file.unlock().map_err_into(HookErrorKind::HookExecutionError))
    }

}

#[derive(PartialEq, Eq, Debug, Clone)]
pub enum Action {
    Lock,
    Unlock
}

fn action_to_str(a: &Action) -> &'static str {
    match *a {
        Action::Lock   => "lock",
        Action::Unlock => "unlock",
    }
}

#[derive(Debug, Clone)]
pub struct FlockUpdateHook {
    action: Action,
}

impl FlockUpdateHook {

    pub fn new(action: Action) -> FlockUpdateHook {
        FlockUpdateHook {
            action: action,
        }
    }

}

impl Hook for FlockUpdateHook {

    fn name(&self) -> &'static str {
        "stdhook_flock_update"
    }

    fn set_config(&mut self, _: &Value) {
        () // We are not configurable here.
    }

}

impl HookDataAccessorProvider for FlockUpdateHook {

    fn accessor(&self) -> HDA {
        HDA::StoreIdAccess(self)
    }

}

impl StoreIdAccessor for FlockUpdateHook {

    fn access(&self, id: &StoreId) -> HookResult<()> {
        debug!("[FLOCK HOOK][{}] {:?}", action_to_str(&self.action), id);
        Ok(())
    }

}

impl MutableHookDataAccessor for FlockUpdateHook {

    fn access_mut(&self, fle: &mut FileLockEntry) -> HookResult<()> {
        debug!("[FLOCK HOOK][{}] {:?}", action_to_str(&self.action), fle.get_location());
        fle.lock()
            .map_err(|e| HookError::new(HookErrorKind::HookExecutionError, Some(Box::new(e))))
            .map(|_| ())
    }

}

impl NonMutableHookDataAccessor for FlockUpdateHook {

    fn access(&self, fle: &FileLockEntry) -> HookResult<()> {
        debug!("[FLOCK HOOK][{}] {:?}", action_to_str(&self.action), fle.get_location());
        fle.unlock()
            .map_err(|e| HookError::new(HookErrorKind::HookExecutionError, Some(Box::new(e))))
            .map(|_| ())
    }

}

<commit_msg>Have and use internal error types for better error messages<commit_after>use toml::Value;

use fs2::FileExt;

use libimagstore::hook::Hook;
use libimagstore::hook::accessor::HookDataAccessor as HDA;
use libimagstore::hook::accessor::HookDataAccessorProvider;
use libimagstore::hook::accessor::StoreIdAccessor;
use libimagstore::hook::accessor::MutableHookDataAccessor;
use libimagstore::hook::accessor::NonMutableHookDataAccessor;
use libimagstore::hook::result::HookResult;
use libimagstore::hook::error::{HookError, HookErrorKind};
use libimagstore::storeid::StoreId;
use libimagstore::store::FileLockEntry;
use libimagstore::store::Entry;

mod error {
    generate_error_imports!();
    generate_error_types!(FlockError, FlockErrorKind,
        IOError                    => "IO Error",
        StoreIdPathBufConvertError => "Error while converting StoreId to PathBuf",
        FileOpenError              => "Error on File::open()",
        LockError                  => "Error while lock()ing",
        UnlockError                => "Error while unlock()ing"
    );
}
use self::error::FlockError as FE;
use self::error::FlockErrorKind as FEK;
use self::error::MapErrInto;

trait EntryFlock {
    fn lock(&self) -> Result<(), FE>;
    fn unlock(&self) -> Result<(), FE>;
}

impl EntryFlock for Entry {

    fn lock(&self) -> Result<(), FE> {
        use std::fs::File;

        self.get_location()
            .clone()
            .into_pathbuf()
            .map_err_into(FEK::StoreIdPathBufConvertError)
            .and_then(|loc| {
                File::open(loc)
                    .map_err_into(FEK::FileOpenError)
                    .map_err_into(FEK::IOError)
            })
            .and_then(|file| {
                file.lock_exclusive()
                    .map_err_into(FEK::LockError)
                    .map_err_into(FEK::IOError)
            })
    }

    fn unlock(&self) -> Result<(), FE> {
        use std::fs::File;

        self.get_location()
            .clone()
            .into_pathbuf()
            .map_err_into(FEK::StoreIdPathBufConvertError)
            .and_then(|loc| {
                File::open(loc)
                    .map_err_into(FEK::FileOpenError)
                    .map_err_into(FEK::IOError)
            })
            .and_then(|file| {
                file.unlock()
                    .map_err_into(FEK::UnlockError)
                    .map_err_into(FEK::LockError)
                    .map_err_into(FEK::IOError)
            })
    }

}

#[derive(PartialEq, Eq, Debug, Clone)]
pub enum Action {
    Lock,
    Unlock
}

fn action_to_str(a: &Action) -> &'static str {
    match *a {
        Action::Lock   => "lock",
        Action::Unlock => "unlock",
    }
}

#[derive(Debug, Clone)]
pub struct FlockUpdateHook {
    action: Action,
}

impl FlockUpdateHook {

    pub fn new(action: Action) -> FlockUpdateHook {
        FlockUpdateHook {
            action: action,
        }
    }

}

impl Hook for FlockUpdateHook {

    fn name(&self) -> &'static str {
        "stdhook_flock_update"
    }

    fn set_config(&mut self, _: &Value) {
        () // We are not configurable here.
    }

}

impl HookDataAccessorProvider for FlockUpdateHook {

    fn accessor(&self) -> HDA {
        HDA::StoreIdAccess(self)
    }

}

impl StoreIdAccessor for FlockUpdateHook {

    fn access(&self, id: &StoreId) -> HookResult<()> {
        debug!("[FLOCK HOOK][{}] {:?}", action_to_str(&self.action), id);
        Ok(())
    }

}

impl MutableHookDataAccessor for FlockUpdateHook {

    fn access_mut(&self, fle: &mut FileLockEntry) -> HookResult<()> {
        debug!("[FLOCK HOOK][{}] {:?}", action_to_str(&self.action), fle.get_location());
        fle.lock()
            .map_err(|e| HookError::new(HookErrorKind::HookExecutionError, Some(Box::new(e))))
            .map(|_| ())
    }

}

impl NonMutableHookDataAccessor for FlockUpdateHook {

    fn access(&self, fle: &FileLockEntry) -> HookResult<()> {
        debug!("[FLOCK HOOK][{}] {:?}", action_to_str(&self.action), fle.get_location());
        fle.unlock()
            .map_err(|e| HookError::new(HookErrorKind::HookExecutionError, Some(Box::new(e))))
            .map(|_| ())
    }

}

<|endoftext|>


<commit_before>use crate::ui::{UIWriter,
                UI};
use habitat_api_client as api_client;
use habitat_core as hcore;
use lazy_static::lazy_static;
use std::{collections::HashMap,
          env,
          ffi::OsStr,
          iter::FromIterator};

extern crate json;
#[macro_use]
extern crate log;
#[macro_use]
extern crate lazy_static;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[cfg_attr(test, macro_use)]
extern crate serde_json;
#[cfg(windows)]
extern crate winapi;

pub use self::error::{Error,
                      Result};

pub mod cli;
pub mod command;
pub mod error;
pub mod locked_env_var;
pub mod output;
pub mod package_graph;
pub mod templating;
pub mod types;
pub mod ui;
pub mod util;

lazy_static::lazy_static! {
    pub static ref PROGRAM_NAME: String = {
        match env::current_exe() {
            Ok(path) => path.file_stem().and_then(OsStr::to_str).unwrap().to_string(),
            Err(e) => {
                error!("Error getting path of current_exe: {}", e);
                String::from("hab-?")
            }
        }
    };
}

// TODO (CM): It would be nice to come up with a way to more
// programmatically manage these flags. It's a bit of a pain to define
// the flag, and then define the environment variables
// separately. Nothing statically guarantees that you've specified an
// variable for a flag.

// TODO (CM): It'd be great to have a built-in way to document them,
// too.

// TODO (CM): Part of that documentation might be *when* a flag was
// added. In general, long-lived flags are a code-smell.

// TODO (CM): It may also be useful to break out features by area of
// concern. We can have any number of bitflags-generated structs.

bitflags::bitflags! {
    /// All the feature flags that are recogized by Habitat.
    ///
    /// In general, feature flags are enabled by setting the corresponding
    /// environment variable.
    ///
    /// Your binary should call `FeatureFlag::from_env` to get a set
    /// of flags to use.
    ///
    /// To add a new feature flag, you will need to add the bit mask
    /// constant here, as well as a mapping from the feature to the
    /// environment variable to which it corresponds in the `ENV_VARS`
    /// map below.
    pub struct FeatureFlag: u32 {
        const LIST               = 0b0000_0000_0001;
        const TEST_EXIT          = 0b0000_0000_0010;
        const TEST_BOOT_FAIL     = 0b0000_0000_0100;
        const REDACT_HTTP        = 0b0000_0000_1000;
        const IGNORE_SIGNALS     = 0b0000_0001_0000;
        const OFFLINE_INSTALL    = 0b0000_0100_0000;
        const IGNORE_LOCAL       = 0b0000_1000_0000;
        const EVENT_STREAM       = 0b0001_0000_0000;
        const TRIGGER_ELECTION   = 0b0010_0000_0000;
        const CONFIGURE_SHUTDOWN = 0b0100_0000_0000;
    }
}

lazy_static! {
    static ref ENV_VARS: HashMap<FeatureFlag, &'static str> = {
        let mapping = vec![(FeatureFlag::LIST, "HAB_FEAT_LIST"),
                           (FeatureFlag::TEST_EXIT, "HAB_FEAT_TEST_EXIT"),
                           (FeatureFlag::TEST_BOOT_FAIL, "HAB_FEAT_BOOT_FAIL"),
                           (FeatureFlag::REDACT_HTTP, "HAB_FEAT_REDACT_HTTP"),
                           (FeatureFlag::IGNORE_SIGNALS, "HAB_FEAT_IGNORE_SIGNALS"),
                           (FeatureFlag::OFFLINE_INSTALL, "HAB_FEAT_OFFLINE_INSTALL"),
                           (FeatureFlag::IGNORE_LOCAL, "HAB_FEAT_IGNORE_LOCAL"),
                           (FeatureFlag::EVENT_STREAM, "HAB_FEAT_EVENT_STREAM"),
                           (FeatureFlag::TRIGGER_ELECTION, "HAB_FEAT_TRIGGER_ELECTION"),
                           (FeatureFlag::CONFIGURE_SHUTDOWN, "HAB_FEAT_CONFIGURE_SHUTDOWN")];
        HashMap::from_iter(mapping)
    };
}

impl FeatureFlag {
    /// If the environment variable for a flag is set to _anything_ but
    /// the empty string, it is activated.
    pub fn from_env(ui: &mut UI) -> Self {
        let mut flags = FeatureFlag::empty();

        for (feature, env_var) in ENV_VARS.iter() {
            if let Some(val) = env::var_os(env_var) {
                if !val.is_empty() {
                    flags.insert(*feature);
                    ui.warn(&format!("Enabling feature: {:?}", feature))
                      .unwrap();
                }
            }
        }

        // TODO (CM): Once the other TODOs above are done (especially the
        // documentation bits), it would be nice to extract this logic
        // into an actual discoverable CLI subcommand; it's a little weird
        // that you have to know how to enable a feature flag before you
        // can even find out that there *are* feature flags to enable.
        //
        // There's no reason why "list feature flags" should itself be a
        // feature-flag.
        if flags.contains(FeatureFlag::LIST) {
            ui.warn("Listing feature flags environment variables:")
              .unwrap();
            for (feature, env_var) in ENV_VARS.iter() {
                ui.warn(&format!("  * {:?}: {}={:?}",
                                 feature,
                                 env_var,
                                 env::var_os(env_var).unwrap_or_default()))
                  .unwrap();
            }
        }

        flags
    }
}

pub mod sync {
    use std::{collections::HashMap,
              sync::Mutex,
              thread::{self,
                       ThreadId},
              time::{Duration,
                     Instant}};

    type NameAndLastHeartbeat = (Option<String>, Instant);
    type HeartbeatMap = HashMap<ThreadId, NameAndLastHeartbeat>;
    lazy_static::lazy_static! {
        static ref THREAD_HEARTBEATS: Mutex<HeartbeatMap> = Default::default();
    }

    habitat_core::env_config_duration!(ThreadAliveThreshold,
                                       HAB_THREAD_ALIVE_THRESHOLD_SECS,
                                       Duration::from_secs(5 * 60));
    habitat_core::env_config_duration!(ThreadAliveCheckDelay,
                                       HAB_THREAD_ALIVE_CHECK_DELAY_SECS,
                                       Duration::from_secs(60));

    /// Call periodically from a thread which has a work loop to indicate that the thread is
    /// still alive and processing its loop. If this function is not called for more than
    /// `ThreadAliveThreshold`, it will be error logged as a likely deadlock.
    pub fn mark_thread_alive() {
        mark_thread_alive_impl(&mut THREAD_HEARTBEATS.lock()
                                                     .expect("THREAD_HEARTBEATS poisoned"));
    }

    fn mark_thread_alive_impl(heartbeats: &mut HeartbeatMap) {
        let thread = thread::current();
        heartbeats.insert(thread.id(),
                          (thread.name().map(str::to_string), Instant::now()));
    }

    /// Call once per binary to start the thread which will check that all the threads that
    /// call `mark_thread_alive` continue to do so.
    pub fn spawn_thread_alive_checker() {
        thread::Builder::new().name("thread-alive-check".to_string())
                              .spawn(|| {
                                  let delay = ThreadAliveCheckDelay::configured_value().into();
                                  let threshold = ThreadAliveThreshold::configured_value().into();
                                  loop {
                                      check_thread_heartbeats(threshold);
                                      thread::sleep(delay);
                                  }
                              })
                              .expect("Error spawning thread alive checker");
    }

    fn check_thread_heartbeats(threshold: Duration) {
        let heartbeats = &THREAD_HEARTBEATS.lock()
                                           .expect("THREAD_HEARTBEATS poisoned");
        for (name, last_heartbeat) in threads_missing_heartbeat(heartbeats, threshold) {
            error!("No heartbeat from {} in {} seconds; deadlock likely",
                   name.unwrap_or_else(|| { "unnamed thread".to_string() }),
                   last_heartbeat.elapsed().as_secs());
        }
    }

    fn threads_missing_heartbeat(heartbeats: &HeartbeatMap,
                                 threshold: Duration)
                                 -> Vec<NameAndLastHeartbeat> {
        heartbeats.iter()
                  .filter_map(|(thread_id, (thread_name, last_heartbeat))| {
                      let time_since_last_heartbeat = last_heartbeat.elapsed();
                      trace!("{:?} {:?} last heartbeat: {:?} ago",
                             thread_id,
                             thread_name,
                             time_since_last_heartbeat);
                      if time_since_last_heartbeat < threshold {
                          None
                      } else {
                          Some((thread_name.clone(), *last_heartbeat))
                      }
                  })
                  .collect::<Vec<_>>()
    }

    #[cfg(test)]
    mod test {
        use super::*;
        use std::sync::{atomic::{AtomicBool,
                                 Ordering},
                        Arc};

        const TEST_THRESHOLD: Duration = Duration::from_millis(10);

        #[test]
        fn no_tracking_without_mark_thread_alive() {
            let heartbeats = HashMap::new();
            thread::spawn(|| {}).join().unwrap();
            thread::sleep(TEST_THRESHOLD * 2);
            assert!(threads_missing_heartbeat(&heartbeats, TEST_THRESHOLD).is_empty());
        }

        #[test]
        fn one_dead_thread() {
            lazy_static! {
                static ref HEARTBEATS: Mutex<HeartbeatMap> = Default::default();
            }
            thread::spawn(move || mark_thread_alive_impl(&mut HEARTBEATS.lock().unwrap())).join()
                                                                                .unwrap();
            thread::sleep(TEST_THRESHOLD * 2);
            assert_eq!(threads_missing_heartbeat(&HEARTBEATS.lock().unwrap(), TEST_THRESHOLD).len(),
                       1);
        }

        #[test]
        fn one_dead_one_alive() {
            lazy_static! {
                static ref HEARTBEATS: Mutex<HeartbeatMap> = Default::default();
            }

            let dead_thread_name = "expected-dead".to_string();
            let test_done: Arc<AtomicBool> = Default::default();
            let test_done2 = Arc::clone(&test_done);

            thread::Builder::new().name(dead_thread_name.clone())
                                  .spawn(move || {
                                      mark_thread_alive_impl(&mut HEARTBEATS.lock().unwrap())
                                  })
                                  .unwrap()
                                  .join()
                                  .unwrap();
            thread::spawn(move || {
                while !test_done2.load(Ordering::Relaxed) {
                    mark_thread_alive_impl(&mut HEARTBEATS.lock().unwrap());
                    thread::sleep(TEST_THRESHOLD / 2)
                }
            });

            thread::sleep(TEST_THRESHOLD * 2);

            let dead_thread_names = threads_missing_heartbeat(&HEARTBEATS.lock().unwrap(),
                                                              TEST_THRESHOLD).iter()
                                                                             .map(|(name, _)| name.clone())
                                                                             .collect::<Vec<_>>();
            assert_eq!(dead_thread_names, vec![Some(dead_thread_name)]);
            test_done.store(true, Ordering::Relaxed);
        }
    }

    #[cfg(feature = "deadlock_detection")]
    mod deadlock_detection {
        use super::*;
        use parking_lot::deadlock;
        use std::{sync::Once,
                  thread};

        static INIT: Once = Once::new();

        pub fn init() { INIT.call_once(spawn_deadlock_detector_thread); }

        fn spawn_deadlock_detector_thread() {
            thread::spawn(move || {
                loop {
                    thread::sleep(Duration::from_secs(10));
                    let deadlocks = deadlock::check_deadlock();
                    if deadlocks.is_empty() {
                        continue;
                    }

                    println!("{} deadlocks detected", deadlocks.len());
                    for (i, threads) in deadlocks.iter().enumerate() {
                        println!("Deadlock #{}", i);
                        for t in threads {
                            println!("Thread Id {:#?}", t.thread_id());
                            println!("{:#?}", t.backtrace());
                        }
                    }

                    // Unfortunately, we can't do anything to resolve the deadlock and
                    // continue, so we have to abort the whole process
                    std::process::exit(1);
                }
            });
        }
    }

    #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
    type InnerLock<T> = parking_lot::RwLock<T>;
    #[cfg(feature = "lock_as_mutex")]
    type InnerLock<T> = parking_lot::Mutex<T>;

    #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
    pub type ReadGuard<'a, T> = parking_lot::RwLockReadGuard<'a, T>;
    #[cfg(feature = "lock_as_mutex")]
    pub type ReadGuard<'a, T> = parking_lot::MutexGuard<'a, T>;

    #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
    pub type WriteGuard<'a, T> = parking_lot::RwLockWriteGuard<'a, T>;
    #[cfg(feature = "lock_as_mutex")]
    pub type WriteGuard<'a, T> = parking_lot::MutexGuard<'a, T>;

    /// A lock which provides the interface of a read/write lock, but which has the option to
    /// internally use either a RwLock or a Mutex in order to make it easier to expose erroneous
    /// recursive locking in tests while still using an RwLock in production to avoid deadlocking
    /// as much as possible.
    #[derive(Debug)]
    pub struct Lock<T> {
        inner: InnerLock<T>,
    }

    impl<T> Lock<T> {
        pub fn new(val: T) -> Self {
            #[cfg(feature = "lock_as_mutex")]
            println!("Lock::new is using Mutex to help find recursive locking");

            #[cfg(feature = "deadlock_detection")]
            deadlock_detection::init();

            Self { inner: InnerLock::new(val), }
        }

        /// This acquires a read lock and will not deadlock if the same thread tries to acquire
        /// the lock recursively. However, it may result in writer starvation. Once we are confident
        /// that all recursive locking has been eliminated, we may replace this implementation
        /// and try_read_for (or add an additional methods) to provide fair locking for readers
        /// and writers.
        ///
        /// See https://github.com/habitat-sh/habitat/issues/6435
        pub fn read(&self) -> ReadGuard<T> {
            #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
            {
                self.inner.read_recursive()
            }
            #[cfg(feature = "lock_as_mutex")]
            {
                self.inner.lock()
            }
        }

        pub fn try_read_for(&self, timeout: Duration) -> Option<ReadGuard<T>> {
            #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
            {
                self.inner.try_read_recursive_for(timeout)
            }
            #[cfg(feature = "lock_as_mutex")]
            {
                self.inner.try_lock_for(timeout)
            }
        }

        pub fn write(&self) -> WriteGuard<T> {
            #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
            {
                self.inner.write()
            }
            #[cfg(feature = "lock_as_mutex")]
            {
                self.inner.lock()
            }
        }
    }
}
<commit_msg>Move deadlock error to trace - stop spam<commit_after>use crate::ui::{UIWriter,
                UI};
use habitat_api_client as api_client;
use habitat_core as hcore;
use lazy_static::lazy_static;
use std::{collections::HashMap,
          env,
          ffi::OsStr,
          iter::FromIterator};

extern crate json;
#[macro_use]
extern crate log;
#[macro_use]
extern crate lazy_static;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[cfg_attr(test, macro_use)]
extern crate serde_json;
#[cfg(windows)]
extern crate winapi;

pub use self::error::{Error,
                      Result};

pub mod cli;
pub mod command;
pub mod error;
pub mod locked_env_var;
pub mod output;
pub mod package_graph;
pub mod templating;
pub mod types;
pub mod ui;
pub mod util;

lazy_static::lazy_static! {
    pub static ref PROGRAM_NAME: String = {
        match env::current_exe() {
            Ok(path) => path.file_stem().and_then(OsStr::to_str).unwrap().to_string(),
            Err(e) => {
                error!("Error getting path of current_exe: {}", e);
                String::from("hab-?")
            }
        }
    };
}

// TODO (CM): It would be nice to come up with a way to more
// programmatically manage these flags. It's a bit of a pain to define
// the flag, and then define the environment variables
// separately. Nothing statically guarantees that you've specified an
// variable for a flag.

// TODO (CM): It'd be great to have a built-in way to document them,
// too.

// TODO (CM): Part of that documentation might be *when* a flag was
// added. In general, long-lived flags are a code-smell.

// TODO (CM): It may also be useful to break out features by area of
// concern. We can have any number of bitflags-generated structs.

bitflags::bitflags! {
    /// All the feature flags that are recogized by Habitat.
    ///
    /// In general, feature flags are enabled by setting the corresponding
    /// environment variable.
    ///
    /// Your binary should call `FeatureFlag::from_env` to get a set
    /// of flags to use.
    ///
    /// To add a new feature flag, you will need to add the bit mask
    /// constant here, as well as a mapping from the feature to the
    /// environment variable to which it corresponds in the `ENV_VARS`
    /// map below.
    pub struct FeatureFlag: u32 {
        const LIST               = 0b0000_0000_0001;
        const TEST_EXIT          = 0b0000_0000_0010;
        const TEST_BOOT_FAIL     = 0b0000_0000_0100;
        const REDACT_HTTP        = 0b0000_0000_1000;
        const IGNORE_SIGNALS     = 0b0000_0001_0000;
        const OFFLINE_INSTALL    = 0b0000_0100_0000;
        const IGNORE_LOCAL       = 0b0000_1000_0000;
        const EVENT_STREAM       = 0b0001_0000_0000;
        const TRIGGER_ELECTION   = 0b0010_0000_0000;
        const CONFIGURE_SHUTDOWN = 0b0100_0000_0000;
    }
}

lazy_static! {
    static ref ENV_VARS: HashMap<FeatureFlag, &'static str> = {
        let mapping = vec![(FeatureFlag::LIST, "HAB_FEAT_LIST"),
                           (FeatureFlag::TEST_EXIT, "HAB_FEAT_TEST_EXIT"),
                           (FeatureFlag::TEST_BOOT_FAIL, "HAB_FEAT_BOOT_FAIL"),
                           (FeatureFlag::REDACT_HTTP, "HAB_FEAT_REDACT_HTTP"),
                           (FeatureFlag::IGNORE_SIGNALS, "HAB_FEAT_IGNORE_SIGNALS"),
                           (FeatureFlag::OFFLINE_INSTALL, "HAB_FEAT_OFFLINE_INSTALL"),
                           (FeatureFlag::IGNORE_LOCAL, "HAB_FEAT_IGNORE_LOCAL"),
                           (FeatureFlag::EVENT_STREAM, "HAB_FEAT_EVENT_STREAM"),
                           (FeatureFlag::TRIGGER_ELECTION, "HAB_FEAT_TRIGGER_ELECTION"),
                           (FeatureFlag::CONFIGURE_SHUTDOWN, "HAB_FEAT_CONFIGURE_SHUTDOWN")];
        HashMap::from_iter(mapping)
    };
}

impl FeatureFlag {
    /// If the environment variable for a flag is set to _anything_ but
    /// the empty string, it is activated.
    pub fn from_env(ui: &mut UI) -> Self {
        let mut flags = FeatureFlag::empty();

        for (feature, env_var) in ENV_VARS.iter() {
            if let Some(val) = env::var_os(env_var) {
                if !val.is_empty() {
                    flags.insert(*feature);
                    ui.warn(&format!("Enabling feature: {:?}", feature))
                      .unwrap();
                }
            }
        }

        // TODO (CM): Once the other TODOs above are done (especially the
        // documentation bits), it would be nice to extract this logic
        // into an actual discoverable CLI subcommand; it's a little weird
        // that you have to know how to enable a feature flag before you
        // can even find out that there *are* feature flags to enable.
        //
        // There's no reason why "list feature flags" should itself be a
        // feature-flag.
        if flags.contains(FeatureFlag::LIST) {
            ui.warn("Listing feature flags environment variables:")
              .unwrap();
            for (feature, env_var) in ENV_VARS.iter() {
                ui.warn(&format!("  * {:?}: {}={:?}",
                                 feature,
                                 env_var,
                                 env::var_os(env_var).unwrap_or_default()))
                  .unwrap();
            }
        }

        flags
    }
}

pub mod sync {
    use std::{collections::HashMap,
              sync::Mutex,
              thread::{self,
                       ThreadId},
              time::{Duration,
                     Instant}};

    type NameAndLastHeartbeat = (Option<String>, Instant);
    type HeartbeatMap = HashMap<ThreadId, NameAndLastHeartbeat>;
    lazy_static::lazy_static! {
        static ref THREAD_HEARTBEATS: Mutex<HeartbeatMap> = Default::default();
    }

    habitat_core::env_config_duration!(ThreadAliveThreshold,
                                       HAB_THREAD_ALIVE_THRESHOLD_SECS,
                                       Duration::from_secs(5 * 60));
    habitat_core::env_config_duration!(ThreadAliveCheckDelay,
                                       HAB_THREAD_ALIVE_CHECK_DELAY_SECS,
                                       Duration::from_secs(60));

    /// Call periodically from a thread which has a work loop to indicate that the thread is
    /// still alive and processing its loop. If this function is not called for more than
    /// `ThreadAliveThreshold`, it will be error logged as a likely deadlock.
    pub fn mark_thread_alive() {
        mark_thread_alive_impl(&mut THREAD_HEARTBEATS.lock()
                                                     .expect("THREAD_HEARTBEATS poisoned"));
    }

    fn mark_thread_alive_impl(heartbeats: &mut HeartbeatMap) {
        let thread = thread::current();
        heartbeats.insert(thread.id(),
                          (thread.name().map(str::to_string), Instant::now()));
    }

    /// Call once per binary to start the thread which will check that all the threads that
    /// call `mark_thread_alive` continue to do so.
    pub fn spawn_thread_alive_checker() {
        thread::Builder::new().name("thread-alive-check".to_string())
                              .spawn(|| {
                                  let delay = ThreadAliveCheckDelay::configured_value().into();
                                  let threshold = ThreadAliveThreshold::configured_value().into();
                                  loop {
                                      check_thread_heartbeats(threshold);
                                      thread::sleep(delay);
                                  }
                              })
                              .expect("Error spawning thread alive checker");
    }

    fn check_thread_heartbeats(threshold: Duration) {
        let heartbeats = &THREAD_HEARTBEATS.lock()
                                           .expect("THREAD_HEARTBEATS poisoned");
        for (name, last_heartbeat) in threads_missing_heartbeat(heartbeats, threshold) {
            warn!("No heartbeat from {} in {} seconds; deadlock likely",
                   name.unwrap_or_else(|| { "unnamed thread".to_string() }),
                   last_heartbeat.elapsed().as_secs());
        }
    }

    fn threads_missing_heartbeat(heartbeats: &HeartbeatMap,
                                 threshold: Duration)
                                 -> Vec<NameAndLastHeartbeat> {
        heartbeats.iter()
                  .filter_map(|(thread_id, (thread_name, last_heartbeat))| {
                      let time_since_last_heartbeat = last_heartbeat.elapsed();
                      trace!("{:?} {:?} last heartbeat: {:?} ago",
                             thread_id,
                             thread_name,
                             time_since_last_heartbeat);
                      if time_since_last_heartbeat < threshold {
                          None
                      } else {
                          Some((thread_name.clone(), *last_heartbeat))
                      }
                  })
                  .collect::<Vec<_>>()
    }

    #[cfg(test)]
    mod test {
        use super::*;
        use std::sync::{atomic::{AtomicBool,
                                 Ordering},
                        Arc};

        const TEST_THRESHOLD: Duration = Duration::from_millis(10);

        #[test]
        fn no_tracking_without_mark_thread_alive() {
            let heartbeats = HashMap::new();
            thread::spawn(|| {}).join().unwrap();
            thread::sleep(TEST_THRESHOLD * 2);
            assert!(threads_missing_heartbeat(&heartbeats, TEST_THRESHOLD).is_empty());
        }

        #[test]
        fn one_dead_thread() {
            lazy_static! {
                static ref HEARTBEATS: Mutex<HeartbeatMap> = Default::default();
            }
            thread::spawn(move || mark_thread_alive_impl(&mut HEARTBEATS.lock().unwrap())).join()
                                                                                .unwrap();
            thread::sleep(TEST_THRESHOLD * 2);
            assert_eq!(threads_missing_heartbeat(&HEARTBEATS.lock().unwrap(), TEST_THRESHOLD).len(),
                       1);
        }

        #[test]
        fn one_dead_one_alive() {
            lazy_static! {
                static ref HEARTBEATS: Mutex<HeartbeatMap> = Default::default();
            }

            let dead_thread_name = "expected-dead".to_string();
            let test_done: Arc<AtomicBool> = Default::default();
            let test_done2 = Arc::clone(&test_done);

            thread::Builder::new().name(dead_thread_name.clone())
                                  .spawn(move || {
                                      mark_thread_alive_impl(&mut HEARTBEATS.lock().unwrap())
                                  })
                                  .unwrap()
                                  .join()
                                  .unwrap();
            thread::spawn(move || {
                while !test_done2.load(Ordering::Relaxed) {
                    mark_thread_alive_impl(&mut HEARTBEATS.lock().unwrap());
                    thread::sleep(TEST_THRESHOLD / 2)
                }
            });

            thread::sleep(TEST_THRESHOLD * 2);

            let dead_thread_names = threads_missing_heartbeat(&HEARTBEATS.lock().unwrap(),
                                                              TEST_THRESHOLD).iter()
                                                                             .map(|(name, _)| name.clone())
                                                                             .collect::<Vec<_>>();
            assert_eq!(dead_thread_names, vec![Some(dead_thread_name)]);
            test_done.store(true, Ordering::Relaxed);
        }
    }

    #[cfg(feature = "deadlock_detection")]
    mod deadlock_detection {
        use super::*;
        use parking_lot::deadlock;
        use std::{sync::Once,
                  thread};

        static INIT: Once = Once::new();

        pub fn init() { INIT.call_once(spawn_deadlock_detector_thread); }

        fn spawn_deadlock_detector_thread() {
            thread::spawn(move || {
                loop {
                    thread::sleep(Duration::from_secs(10));
                    let deadlocks = deadlock::check_deadlock();
                    if deadlocks.is_empty() {
                        continue;
                    }

                    println!("{} deadlocks detected", deadlocks.len());
                    for (i, threads) in deadlocks.iter().enumerate() {
                        println!("Deadlock #{}", i);
                        for t in threads {
                            println!("Thread Id {:#?}", t.thread_id());
                            println!("{:#?}", t.backtrace());
                        }
                    }

                    // Unfortunately, we can't do anything to resolve the deadlock and
                    // continue, so we have to abort the whole process
                    std::process::exit(1);
                }
            });
        }
    }

    #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
    type InnerLock<T> = parking_lot::RwLock<T>;
    #[cfg(feature = "lock_as_mutex")]
    type InnerLock<T> = parking_lot::Mutex<T>;

    #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
    pub type ReadGuard<'a, T> = parking_lot::RwLockReadGuard<'a, T>;
    #[cfg(feature = "lock_as_mutex")]
    pub type ReadGuard<'a, T> = parking_lot::MutexGuard<'a, T>;

    #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
    pub type WriteGuard<'a, T> = parking_lot::RwLockWriteGuard<'a, T>;
    #[cfg(feature = "lock_as_mutex")]
    pub type WriteGuard<'a, T> = parking_lot::MutexGuard<'a, T>;

    /// A lock which provides the interface of a read/write lock, but which has the option to
    /// internally use either a RwLock or a Mutex in order to make it easier to expose erroneous
    /// recursive locking in tests while still using an RwLock in production to avoid deadlocking
    /// as much as possible.
    #[derive(Debug)]
    pub struct Lock<T> {
        inner: InnerLock<T>,
    }

    impl<T> Lock<T> {
        pub fn new(val: T) -> Self {
            #[cfg(feature = "lock_as_mutex")]
            println!("Lock::new is using Mutex to help find recursive locking");

            #[cfg(feature = "deadlock_detection")]
            deadlock_detection::init();

            Self { inner: InnerLock::new(val), }
        }

        /// This acquires a read lock and will not deadlock if the same thread tries to acquire
        /// the lock recursively. However, it may result in writer starvation. Once we are confident
        /// that all recursive locking has been eliminated, we may replace this implementation
        /// and try_read_for (or add an additional methods) to provide fair locking for readers
        /// and writers.
        ///
        /// See https://github.com/habitat-sh/habitat/issues/6435
        pub fn read(&self) -> ReadGuard<T> {
            #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
            {
                self.inner.read_recursive()
            }
            #[cfg(feature = "lock_as_mutex")]
            {
                self.inner.lock()
            }
        }

        pub fn try_read_for(&self, timeout: Duration) -> Option<ReadGuard<T>> {
            #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
            {
                self.inner.try_read_recursive_for(timeout)
            }
            #[cfg(feature = "lock_as_mutex")]
            {
                self.inner.try_lock_for(timeout)
            }
        }

        pub fn write(&self) -> WriteGuard<T> {
            #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
            {
                self.inner.write()
            }
            #[cfg(feature = "lock_as_mutex")]
            {
                self.inner.lock()
            }
        }
    }
}
<|endoftext|>