<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

extern crate logger;
extern crate sys_util;

use std::io;
use std::result::Result;

use libc::{_exit, c_int, c_void, siginfo_t, SIGBUS, SIGSEGV, SIGSYS};

use logger::{Metric, LOGGER, METRICS};
use sys_util::register_signal_handler;

// The offset of `si_syscall` (offending syscall identifier) within the siginfo structure
// expressed as an `(u)int*`.
// Offset `6` for an `i32` field means that the needed information is located at `6 * sizeof(i32)`.
// See /usr/include/linux/signal.h for the C struct definition.
// See https://github.com/rust-lang/libc/issues/716 for why the offset is different in Rust.
const SI_OFF_SYSCALL: isize = 6;

const SYS_SECCOMP_CODE: i32 = 1;

/// Signal handler for `SIGSYS`.
///
/// Increments the `seccomp.num_faults` metric, logs an error message and terminates the process
/// with a specific exit code.
///
extern "C" fn sigsys_handler(num: c_int, info: *mut siginfo_t, _unused: *mut c_void) {
    // Safe because we're just reading some fields from a supposedly valid argument.
    let si_signo = unsafe { (*info).si_signo };
    let si_code = unsafe { (*info).si_code };

    // Sanity check. The condition should never be true.
    if num != si_signo || num != SIGSYS || si_code != SYS_SECCOMP_CODE as i32 {
        // Safe because we're terminating the process anyway.
        unsafe { _exit(i32::from(super::FC_EXIT_CODE_UNEXPECTED_ERROR)) };
    }

    // Other signals which might do async unsafe things incompatible with the rest of this
    // function are blocked due to the sa_mask used when registering the signal handler.
    let syscall = unsafe { *(info as *const i32).offset(SI_OFF_SYSCALL) as usize };
    METRICS.seccomp.num_faults.inc();
    error!(
        "Shutting down VM after intercepting a bad syscall ({}).",
        syscall
    );
    // Log the metrics before exiting.
    if let Err(e) = LOGGER.log_metrics() {
        error!("Failed to log metrics while stopping: {}", e);
    }

    // Safe because we're terminating the process anyway. We don't actually do anything when
    // running unit tests.
    #[cfg(not(test))]
    unsafe {
        _exit(i32::from(super::FC_EXIT_CODE_BAD_SYSCALL))
    };
}

/// Signal handler for `SIGBUS` and `SIGSEGV`.
///
/// Logs an error message and terminates the process with a specific exit code.
///
extern "C" fn sigbus_sigsegv_handler(num: c_int, info: *mut siginfo_t, _unused: *mut c_void) {
    // Safe because we're just reading some fields from a supposedly valid argument.
    let si_signo = unsafe { (*info).si_signo };
    let si_code = unsafe { (*info).si_code };

    // Sanity check. The condition should never be true.
    if num != si_signo || (num != SIGBUS && num != SIGSEGV) {
        // Safe because we're terminating the process anyway.
        unsafe { _exit(i32::from(super::FC_EXIT_CODE_UNEXPECTED_ERROR)) };
    }

    // Other signals which might do async unsafe things incompatible with the rest of this
    // function are blocked due to the sa_mask used when registering the signal handler.
    error!(
        "Shutting down VM after intercepting signal {}, code {}.",
        si_signo, si_code
    );
    // Log the metrics before exiting.
    if let Err(e) = LOGGER.log_metrics() {
        error!("Failed to log metrics while stopping: {}", e);
    }

    // Safe because we're terminating the process anyway. We don't actually do anything when
    // running unit tests.
    #[cfg(not(test))]
    unsafe {
        _exit(i32::from(match si_signo {
            SIGBUS => super::FC_EXIT_CODE_SIGBUS,
            SIGSEGV => super::FC_EXIT_CODE_SIGSEGV,
            _ => super::FC_EXIT_CODE_UNEXPECTED_ERROR,
        }))
    };
}

/// Registers all the required signal handlers.
///
/// Custom handlers are installed for: `SIGBUS`, `SIGSEGV`, `SIGSYS`.
///
pub fn register_signal_handlers() -> Result<(), io::Error> {
    register_signal_handler(SIGSYS, sigsys_handler)?;
    register_signal_handler(SIGBUS, sigbus_sigsegv_handler)?;
    register_signal_handler(SIGSEGV, sigbus_sigsegv_handler)?;
    Ok(())
}

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

    use std::mem;
    use std::process;

    use libc::{cpu_set_t, syscall};

    use seccomp::{allow_syscall, SeccompAction, SeccompFilter};

    // This function is used when running unit tests, so all the unsafes are safe.
    fn cpu_count() -> usize {
        let mut cpuset: cpu_set_t = unsafe { mem::zeroed() };
        unsafe {
            libc::CPU_ZERO(&mut cpuset);
        }
        let ret = unsafe {
            libc::sched_getaffinity(
                0,
                mem::size_of::<cpu_set_t>(),
                &mut cpuset as *mut cpu_set_t,
            )
        };
        assert_eq!(ret, 0);

        let mut num = 0;
        for i in 0..libc::CPU_SETSIZE as usize {
            if unsafe { libc::CPU_ISSET(i, &cpuset) } {
                num += 1;
            }
        }
        num
    }

    #[test]
    #[should_panic(expected = "process alive")]
    fn test_signal_handler() {
        assert!(register_signal_handlers().is_ok());

        let filter = SeccompFilter::new(
            vec![
                allow_syscall(libc::SYS_brk),
                allow_syscall(libc::SYS_exit),
                allow_syscall(libc::SYS_futex),
                allow_syscall(libc::SYS_getpid),
                allow_syscall(libc::SYS_kill),
                allow_syscall(libc::SYS_munmap),
                allow_syscall(libc::SYS_rt_sigprocmask),
                allow_syscall(libc::SYS_rt_sigreturn),
                allow_syscall(libc::SYS_sched_getaffinity),
                allow_syscall(libc::SYS_set_tid_address),
                allow_syscall(libc::SYS_sigaltstack),
                allow_syscall(libc::SYS_write),
            ]
            .into_iter()
            .collect(),
            SeccompAction::Trap,
        )
        .unwrap();

        assert!(filter.apply().is_ok());
        assert_eq!(METRICS.seccomp.num_faults.count(), 0);

        // Call the blacklisted `SYS_mkdir`.
        unsafe { syscall(libc::SYS_mkdirat, "/foo/bar\0") };

        assert!(cpu_count() > 0);

        // Kcov somehow messes with our handler getting the SIGSYS signal when a bad syscall
        // is caught, so the following assertion no longer holds. Ideally, we'd have a surefire
        // way of either preventing this behaviour, or detecting for certain whether this test is
        // run by kcov or not. The best we could do so far is to look at the perceived number of
        // available CPUs. Kcov seems to make a single CPU available to the process running the
        // tests, so we use this as an heuristic to decide if we check the assertion.
        if cpu_count() > 1 {
            // The signal handler should let the program continue during unit tests.
            assert_eq!(METRICS.seccomp.num_faults.count(), 1);
        }

        // Assert that the SIGBUS handler left the process alive.
        unsafe {
            syscall(libc::SYS_kill, process::id(), SIGBUS);
        }
        // Generate a panic to make sure the process is still alive. If the process exited before,
        // the following panic will not be executed and the test will fail.
        panic!("process alive");
    }
}
<commit_msg>vmm: run test_signal_handler in separate thread<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

extern crate logger;
extern crate sys_util;

use std::io;
use std::result::Result;

use libc::{_exit, c_int, c_void, siginfo_t, SIGBUS, SIGSEGV, SIGSYS};

use logger::{Metric, LOGGER, METRICS};
use sys_util::register_signal_handler;

// The offset of `si_syscall` (offending syscall identifier) within the siginfo structure
// expressed as an `(u)int*`.
// Offset `6` for an `i32` field means that the needed information is located at `6 * sizeof(i32)`.
// See /usr/include/linux/signal.h for the C struct definition.
// See https://github.com/rust-lang/libc/issues/716 for why the offset is different in Rust.
const SI_OFF_SYSCALL: isize = 6;

const SYS_SECCOMP_CODE: i32 = 1;

/// Signal handler for `SIGSYS`.
///
/// Increments the `seccomp.num_faults` metric, logs an error message and terminates the process
/// with a specific exit code.
///
extern "C" fn sigsys_handler(num: c_int, info: *mut siginfo_t, _unused: *mut c_void) {
    // Safe because we're just reading some fields from a supposedly valid argument.
    let si_signo = unsafe { (*info).si_signo };
    let si_code = unsafe { (*info).si_code };

    // Sanity check. The condition should never be true.
    if num != si_signo || num != SIGSYS || si_code != SYS_SECCOMP_CODE as i32 {
        // Safe because we're terminating the process anyway.
        unsafe { _exit(i32::from(super::FC_EXIT_CODE_UNEXPECTED_ERROR)) };
    }

    // Other signals which might do async unsafe things incompatible with the rest of this
    // function are blocked due to the sa_mask used when registering the signal handler.
    let syscall = unsafe { *(info as *const i32).offset(SI_OFF_SYSCALL) as usize };
    METRICS.seccomp.num_faults.inc();
    error!(
        "Shutting down VM after intercepting a bad syscall ({}).",
        syscall
    );
    // Log the metrics before exiting.
    if let Err(e) = LOGGER.log_metrics() {
        error!("Failed to log metrics while stopping: {}", e);
    }

    // Safe because we're terminating the process anyway. We don't actually do anything when
    // running unit tests.
    #[cfg(not(test))]
    unsafe {
        _exit(i32::from(super::FC_EXIT_CODE_BAD_SYSCALL))
    };
}

/// Signal handler for `SIGBUS` and `SIGSEGV`.
///
/// Logs an error message and terminates the process with a specific exit code.
///
extern "C" fn sigbus_sigsegv_handler(num: c_int, info: *mut siginfo_t, _unused: *mut c_void) {
    // Safe because we're just reading some fields from a supposedly valid argument.
    let si_signo = unsafe { (*info).si_signo };
    let si_code = unsafe { (*info).si_code };

    // Sanity check. The condition should never be true.
    if num != si_signo || (num != SIGBUS && num != SIGSEGV) {
        // Safe because we're terminating the process anyway.
        unsafe { _exit(i32::from(super::FC_EXIT_CODE_UNEXPECTED_ERROR)) };
    }

    // Other signals which might do async unsafe things incompatible with the rest of this
    // function are blocked due to the sa_mask used when registering the signal handler.
    error!(
        "Shutting down VM after intercepting signal {}, code {}.",
        si_signo, si_code
    );
    // Log the metrics before exiting.
    if let Err(e) = LOGGER.log_metrics() {
        error!("Failed to log metrics while stopping: {}", e);
    }

    // Safe because we're terminating the process anyway. We don't actually do anything when
    // running unit tests.
    #[cfg(not(test))]
    unsafe {
        _exit(i32::from(match si_signo {
            SIGBUS => super::FC_EXIT_CODE_SIGBUS,
            SIGSEGV => super::FC_EXIT_CODE_SIGSEGV,
            _ => super::FC_EXIT_CODE_UNEXPECTED_ERROR,
        }))
    };
}

/// Registers all the required signal handlers.
///
/// Custom handlers are installed for: `SIGBUS`, `SIGSEGV`, `SIGSYS`.
///
pub fn register_signal_handlers() -> Result<(), io::Error> {
    register_signal_handler(SIGSYS, sigsys_handler)?;
    register_signal_handler(SIGBUS, sigbus_sigsegv_handler)?;
    register_signal_handler(SIGSEGV, sigbus_sigsegv_handler)?;
    Ok(())
}

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

    use libc::{cpu_set_t, syscall};
    use std::{mem, process, thread};

    use seccomp::{allow_syscall, SeccompAction, SeccompFilter};

    // This function is used when running unit tests, so all the unsafes are safe.
    fn cpu_count() -> usize {
        let mut cpuset: cpu_set_t = unsafe { mem::zeroed() };
        unsafe {
            libc::CPU_ZERO(&mut cpuset);
        }
        let ret = unsafe {
            libc::sched_getaffinity(
                0,
                mem::size_of::<cpu_set_t>(),
                &mut cpuset as *mut cpu_set_t,
            )
        };
        assert_eq!(ret, 0);

        let mut num = 0;
        for i in 0..libc::CPU_SETSIZE as usize {
            if unsafe { libc::CPU_ISSET(i, &cpuset) } {
                num += 1;
            }
        }
        num
    }

    #[test]
    fn test_signal_handler() {
        let child = thread::spawn(move || {
            assert!(register_signal_handlers().is_ok());

            let filter = SeccompFilter::new(
                vec![
                    allow_syscall(libc::SYS_brk),
                    allow_syscall(libc::SYS_exit),
                    allow_syscall(libc::SYS_futex),
                    allow_syscall(libc::SYS_getpid),
                    allow_syscall(libc::SYS_munmap),
                    allow_syscall(libc::SYS_kill),
                    allow_syscall(libc::SYS_rt_sigprocmask),
                    allow_syscall(libc::SYS_rt_sigreturn),
                    allow_syscall(libc::SYS_sched_getaffinity),
                    allow_syscall(libc::SYS_set_tid_address),
                    allow_syscall(libc::SYS_sigaltstack),
                    allow_syscall(libc::SYS_write),
                ]
                .into_iter()
                .collect(),
                SeccompAction::Trap,
            )
            .unwrap();

            assert!(filter.apply().is_ok());
            assert_eq!(METRICS.seccomp.num_faults.count(), 0);

            // Call the blacklisted `SYS_mkdirat`.
            unsafe { syscall(libc::SYS_mkdirat, "/foo/bar\0") };

            assert!(cpu_count() > 0);

            // Kcov somehow messes with our handler getting the SIGSYS signal when a bad syscall
            // is caught, so the following assertion no longer holds. Ideally, we'd have a surefire
            // way of either preventing this behaviour, or detecting for certain whether this test is
            // run by kcov or not. The best we could do so far is to look at the perceived number of
            // available CPUs. Kcov seems to make a single CPU available to the process running the
            // tests, so we use this as an heuristic to decide if we check the assertion.
            if cpu_count() > 1 {
                // The signal handler should let the program continue during unit tests.
                assert_eq!(METRICS.seccomp.num_faults.count(), 1);
            }

            // Call SIGBUS signal handler. We are not testing that the SIGBUS handler was executed
            // successfully.
            // Tracking issue: https://github.com/firecracker-microvm/firecracker/issues/1141
            unsafe {
                syscall(libc::SYS_kill, process::id(), SIGBUS);
            }
        });
        assert!(child.join().is_ok());
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]
use ffmpeg_sys;
use libc::*;
use std::cmp;
use std::ffi::{ CStr, CString };
use std::ops::Deref;
use std::ptr;

use errors::*;

#[derive(Debug, Clone, Copy)]
pub struct Rational {
    pub numerator: i32,
    pub denominator: i32,
}

#[derive(Clone, Copy)]
pub struct Codec {
    ptr: *mut ffmpeg_sys::AVCodec,
}

unsafe impl Send for Codec {}
unsafe impl Sync for Codec {}

pub struct PixelFormats {
    codec: Codec,
    index: isize,
}

pub struct Context {
    // The Codec which was used to initialize this Context.
    codec: Codec,
    ptr: *mut ffmpeg_sys::AVCodecContext,
}

pub struct OpenContext {
    context: Context,
}

pub struct Frame {
    ptr: *mut ffmpeg_sys::AVFrame,
}

pub struct OutputContext {
    ptr: *mut ffmpeg_sys::AVFormatContext,
}

impl From<(i32, i32)> for Rational {
    fn from((numerator, denominator): (i32, i32)) -> Self {
        Self {
            numerator,
            denominator,
        }
    }
}

impl From<ffmpeg_sys::AVRational> for Rational {
    fn from(rational: ffmpeg_sys::AVRational) -> Self {
        Self {
            numerator: rational.num,
            denominator: rational.den,
        }
    }
}

impl From<Rational> for ffmpeg_sys::AVRational {
    fn from(rational: Rational) -> Self {
        Self {
            num: rational.numerator,
            den: rational.denominator,
        }
    }
}

impl Codec {
    pub fn description(self) -> String {
        unsafe {
            CStr::from_ptr((*self.ptr).long_name).to_string_lossy().into_owned()
        }
    }

    #[inline]
    pub fn kind(self) -> ffmpeg_sys::AVMediaType {
        unsafe {
            (*self.ptr).kind
        }
    }

    #[inline]
    pub fn is_video(self) -> bool {
        self.kind() == ffmpeg_sys::AVMediaType::AVMEDIA_TYPE_VIDEO
    }

    pub fn pixel_formats(self) -> Option<PixelFormats> {
        unsafe {
            if (*self.ptr).pix_fmts.is_null() {
                None
            } else {
                Some(PixelFormats::new(self))
            }
        }
    }

    pub fn context(self) -> Result<Context> {
        let ptr = unsafe {
            ffmpeg_sys::avcodec_alloc_context3(self.ptr)
        };

        ensure!(!ptr.is_null(), "unable to allocate the codec context");

        Ok(Context {
            codec: self,
            ptr,
        })
    }
}

impl PixelFormats {
    fn new(codec: Codec) -> Self {
        Self {
            codec,
            index: 0,
        }
    }
}

impl Iterator for PixelFormats {
    type Item = ffmpeg_sys::AVPixelFormat;

    fn next(&mut self) -> Option<Self::Item> {
        let format = unsafe {
            *(*self.codec.ptr).pix_fmts.offset(self.index)
        };

        if format == ffmpeg_sys::AVPixelFormat::AV_PIX_FMT_NONE {
            None
        } else {
            self.index += 1;
            Some(format)
        }
    }
}

impl Context {
    #[inline]
    pub fn width(&self) -> u32 {
        unsafe {
            cmp::max(0, (*self.ptr).width) as u32
        }
    }

    #[inline]
    pub fn set_width(&mut self, width: u32) {
        unsafe {
            (*self.ptr).width = cmp::min(width, c_int::max_value() as u32) as c_int;
        }
    }

    #[inline]
    pub fn height(&self) -> u32 {
        unsafe {
            cmp::max(0, (*self.ptr).height) as u32
        }
    }

    #[inline]
    pub fn set_height(&mut self, height: u32) {
        unsafe {
            (*self.ptr).height = cmp::min(height, c_int::max_value() as u32) as c_int;
        }
    }

    #[inline]
    pub fn time_base(&self) -> Rational {
        unsafe {
            (*self.ptr).time_base.into()
        }
    }

    #[inline]
    pub fn set_time_base(&mut self, time_base: &Rational) {
        unsafe {
            (*self.ptr).time_base = (*time_base).into();
        }
    }

    #[inline]
    pub fn pixel_format(&self) -> ffmpeg_sys::AVPixelFormat {
        unsafe {
            (*self.ptr).pix_fmt
        }
    }

    #[inline]
    pub fn set_pixel_format(&mut self, pixel_format: ffmpeg_sys::AVPixelFormat) {
        unsafe {
            (*self.ptr).pix_fmt = pixel_format;
        }
    }

    pub fn open(self) -> Result<OpenContext> {
        let rv = unsafe {
            ffmpeg_sys::avcodec_open2(self.ptr, self.codec.ptr, ptr::null_mut())
        };

        ensure!(rv == 0, "error opening context");

        Ok(OpenContext { context: self })
    }
}

impl Drop for Context {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::avcodec_free_context(&mut self.ptr);
        }
    }
}

impl Frame {
    pub fn new(pixel_format: ffmpeg_sys::AVPixelFormat,
               width: u32,
               height: u32) -> Result<Self> {
        let ptr = unsafe {
            ffmpeg_sys::av_frame_alloc()
        };

        ensure!(!ptr.is_null(), "unable to allocate a frame");

        unsafe {
            (*ptr).format = pixel_format as c_int;
            (*ptr).width = cmp::min(width, c_int::max_value() as u32) as c_int;
            (*ptr).height = cmp::min(height, c_int::max_value() as u32) as c_int;
        }

        Ok(Self { ptr })
    }
}

impl Drop for Frame {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::av_frame_free(&mut self.ptr);
        }
    }
}

impl OutputContext {
    pub fn new(filename: &str) -> Result<Self> {
        let filename = CString::new(filename)
            .chain_err(|| "unable to convert filename to a CString")?;

        let mut ptr = ptr::null_mut();

        let rv = unsafe {
            ffmpeg_sys::avformat_alloc_output_context2(&mut ptr,
                                                       ptr::null_mut(),
                                                       ptr::null_mut(),
                                                       filename.as_ptr())
        };

        // TODO: check and report the error code.
        ensure!(rv >= 0, "unable to allocate the output context");

        Ok(Self { ptr })
    }
}

impl Drop for OutputContext {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::avformat_free_context(self.ptr);
        }
    }
}

pub fn initialize() {
    unsafe {
        ffmpeg_sys::av_register_all();
        ffmpeg_sys::avcodec_register_all();
    }
}

pub fn find_encoder_by_name(name: &str) -> Result<Option<Codec>> {
    let name = CString::new(name).chain_err(|| "could not convert name to CString")?;

    let codec = unsafe {
        ffmpeg_sys::avcodec_find_encoder_by_name(name.as_ptr())
    };

    if codec.is_null() {
        Ok(None)
    } else {
        Ok(Some(Codec { ptr: codec }))
    }
}
<commit_msg>Remove unused use<commit_after>#![allow(dead_code)]
use ffmpeg_sys;
use libc::*;
use std::cmp;
use std::ffi::{ CStr, CString };
use std::ptr;

use errors::*;

#[derive(Debug, Clone, Copy)]
pub struct Rational {
    pub numerator: i32,
    pub denominator: i32,
}

#[derive(Clone, Copy)]
pub struct Codec {
    ptr: *mut ffmpeg_sys::AVCodec,
}

unsafe impl Send for Codec {}
unsafe impl Sync for Codec {}

pub struct PixelFormats {
    codec: Codec,
    index: isize,
}

pub struct Context {
    // The Codec which was used to initialize this Context.
    codec: Codec,
    ptr: *mut ffmpeg_sys::AVCodecContext,
}

pub struct OpenContext {
    context: Context,
}

pub struct Frame {
    ptr: *mut ffmpeg_sys::AVFrame,
}

pub struct OutputContext {
    ptr: *mut ffmpeg_sys::AVFormatContext,
}

impl From<(i32, i32)> for Rational {
    fn from((numerator, denominator): (i32, i32)) -> Self {
        Self {
            numerator,
            denominator,
        }
    }
}

impl From<ffmpeg_sys::AVRational> for Rational {
    fn from(rational: ffmpeg_sys::AVRational) -> Self {
        Self {
            numerator: rational.num,
            denominator: rational.den,
        }
    }
}

impl From<Rational> for ffmpeg_sys::AVRational {
    fn from(rational: Rational) -> Self {
        Self {
            num: rational.numerator,
            den: rational.denominator,
        }
    }
}

impl Codec {
    pub fn description(self) -> String {
        unsafe {
            CStr::from_ptr((*self.ptr).long_name).to_string_lossy().into_owned()
        }
    }

    #[inline]
    pub fn kind(self) -> ffmpeg_sys::AVMediaType {
        unsafe {
            (*self.ptr).kind
        }
    }

    #[inline]
    pub fn is_video(self) -> bool {
        self.kind() == ffmpeg_sys::AVMediaType::AVMEDIA_TYPE_VIDEO
    }

    pub fn pixel_formats(self) -> Option<PixelFormats> {
        unsafe {
            if (*self.ptr).pix_fmts.is_null() {
                None
            } else {
                Some(PixelFormats::new(self))
            }
        }
    }

    pub fn context(self) -> Result<Context> {
        let ptr = unsafe {
            ffmpeg_sys::avcodec_alloc_context3(self.ptr)
        };

        ensure!(!ptr.is_null(), "unable to allocate the codec context");

        Ok(Context {
            codec: self,
            ptr,
        })
    }
}

impl PixelFormats {
    fn new(codec: Codec) -> Self {
        Self {
            codec,
            index: 0,
        }
    }
}

impl Iterator for PixelFormats {
    type Item = ffmpeg_sys::AVPixelFormat;

    fn next(&mut self) -> Option<Self::Item> {
        let format = unsafe {
            *(*self.codec.ptr).pix_fmts.offset(self.index)
        };

        if format == ffmpeg_sys::AVPixelFormat::AV_PIX_FMT_NONE {
            None
        } else {
            self.index += 1;
            Some(format)
        }
    }
}

impl Context {
    #[inline]
    pub fn width(&self) -> u32 {
        unsafe {
            cmp::max(0, (*self.ptr).width) as u32
        }
    }

    #[inline]
    pub fn set_width(&mut self, width: u32) {
        unsafe {
            (*self.ptr).width = cmp::min(width, c_int::max_value() as u32) as c_int;
        }
    }

    #[inline]
    pub fn height(&self) -> u32 {
        unsafe {
            cmp::max(0, (*self.ptr).height) as u32
        }
    }

    #[inline]
    pub fn set_height(&mut self, height: u32) {
        unsafe {
            (*self.ptr).height = cmp::min(height, c_int::max_value() as u32) as c_int;
        }
    }

    #[inline]
    pub fn time_base(&self) -> Rational {
        unsafe {
            (*self.ptr).time_base.into()
        }
    }

    #[inline]
    pub fn set_time_base(&mut self, time_base: &Rational) {
        unsafe {
            (*self.ptr).time_base = (*time_base).into();
        }
    }

    #[inline]
    pub fn pixel_format(&self) -> ffmpeg_sys::AVPixelFormat {
        unsafe {
            (*self.ptr).pix_fmt
        }
    }

    #[inline]
    pub fn set_pixel_format(&mut self, pixel_format: ffmpeg_sys::AVPixelFormat) {
        unsafe {
            (*self.ptr).pix_fmt = pixel_format;
        }
    }

    pub fn open(self) -> Result<OpenContext> {
        let rv = unsafe {
            ffmpeg_sys::avcodec_open2(self.ptr, self.codec.ptr, ptr::null_mut())
        };

        ensure!(rv == 0, "error opening context");

        Ok(OpenContext { context: self })
    }
}

impl Drop for Context {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::avcodec_free_context(&mut self.ptr);
        }
    }
}

impl Frame {
    pub fn new(pixel_format: ffmpeg_sys::AVPixelFormat,
               width: u32,
               height: u32) -> Result<Self> {
        let ptr = unsafe {
            ffmpeg_sys::av_frame_alloc()
        };

        ensure!(!ptr.is_null(), "unable to allocate a frame");

        unsafe {
            (*ptr).format = pixel_format as c_int;
            (*ptr).width = cmp::min(width, c_int::max_value() as u32) as c_int;
            (*ptr).height = cmp::min(height, c_int::max_value() as u32) as c_int;
        }

        Ok(Self { ptr })
    }
}

impl Drop for Frame {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::av_frame_free(&mut self.ptr);
        }
    }
}

impl OutputContext {
    pub fn new(filename: &str) -> Result<Self> {
        let filename = CString::new(filename)
            .chain_err(|| "unable to convert filename to a CString")?;

        let mut ptr = ptr::null_mut();

        let rv = unsafe {
            ffmpeg_sys::avformat_alloc_output_context2(&mut ptr,
                                                       ptr::null_mut(),
                                                       ptr::null_mut(),
                                                       filename.as_ptr())
        };

        // TODO: check and report the error code.
        ensure!(rv >= 0, "unable to allocate the output context");

        Ok(Self { ptr })
    }
}

impl Drop for OutputContext {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::avformat_free_context(self.ptr);
        }
    }
}

pub fn initialize() {
    unsafe {
        ffmpeg_sys::av_register_all();
        ffmpeg_sys::avcodec_register_all();
    }
}

pub fn find_encoder_by_name(name: &str) -> Result<Option<Codec>> {
    let name = CString::new(name).chain_err(|| "could not convert name to CString")?;

    let codec = unsafe {
        ffmpeg_sys::avcodec_find_encoder_by_name(name.as_ptr())
    };

    if codec.is_null() {
        Ok(None)
    } else {
        Ok(Some(Codec { ptr: codec }))
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::cell::RefCell;

use conf::{RemoteBranch};
use constants::{CHANGED_KEY,NEW_KEY,STAGED_KEY,CONFLICTS_KEY};

use git2::*;

struct Cache {
    current_branch_name: RefCell<Option<String>>,
    // TODO: Reference can't be cached (can't be cloned)
    //       implement via OID
    // head: RefCell<Option<Oid>>,
}

pub struct Backend {
    cache: Cache,
    pub repo: Repository,
    pub debug: bool,
}

pub struct BranchAheadBehind {
    pub local_branch_name: Option<String>,
    pub remote_branch_name: Option<String>,
    pub remote_name: Option<String>,
    pub ahead: usize,
    pub behind: usize
}

impl BranchAheadBehind {
    fn new(l: Option<String>) -> BranchAheadBehind {
        BranchAheadBehind{ local_branch_name: l, remote_branch_name: None, remote_name: None,
                           ahead: 0, behind: 0 }
    }
}

#[derive(Clone)]
struct RefPair {
    remote_name: String,
    branch_name: String,
    oid: Oid,
}

impl Cache {
    pub fn new() -> Cache {
        Cache{ current_branch_name: RefCell::new(None) }
    }

    fn set_current_branch_name(&self, n: Option<String>) {
        let mut c = self.current_branch_name.borrow_mut();
        *c = n;
    }

    fn is_current_branch_set(&self) -> bool {
        self.current_branch_name.borrow().is_some()
    }

    fn get_current_branch(&self) -> Option<String> {
        self.current_branch_name.borrow().clone()
    }
}

impl Backend {
    pub fn new(repo: Repository, debug: bool) -> Backend {
        Backend{ repo: repo, debug: debug, cache: Cache::new() }
    }

    fn get_head(&self) -> Option<Reference> {
        match self.repo.head() {
            Ok(head) => Some(head),
            Err(e2) => {
                log!(self, "Can't get HEAD: {}", e2);
                match self.repo.find_reference("HEAD") {
                    Ok(x) => {
                        log!(self, "Found HEAD directly: {:?}", x.name());
                        Some(x)
                    },
                    Err(e) => {
                        log!(self, "reference HEAD: {}", e);
                        return None
                    },
                }
            },
        }
    }

    fn get_current_branch_oid(&self) -> Option<Oid> {
        let head = self.get_head();
        match head {
            Some(v) => v.target(),
            None => {
                log!(self, "Failed to find Oid for HEAD.");
                return None
            },
        }
    }

    fn resolve_symbolic_reference<'a>(&self, reference: Option<Reference<'a>>) -> Option<Reference<'a>> {
        let s_reference = match reference {
            Some(v) => {
                match v.resolve() {
                    Ok(y) => {
                        log!(self, "Branch name is {:?}", y.name());
                        y
                    },
                    Err(e) => {
                        log!(self, "Branch name is {:?}, error: {:?}", v.name(), e);
                        v
                    }
                }
            }
            None => {
                log!(self, "No branch name found");
                return None;
            }
        };
        Some(s_reference)
    }

    fn get_branch_name_for_reference(&self, r: Reference) -> Option<String> {
        if let Some(ref_name_string) = r.shorthand() {
            if ref_name_string != "HEAD" {
                let s = ref_name_string.to_string();
                log!(self, "Shorthand for reference is: {}", s);
                return Some(s);
            } else {
                if let Some(ref_name) = r.symbolic_target() {
                    let ref_name_string = ref_name.to_string();
                    log!(self, "shorthand = HEAD, links to: {}", ref_name_string);
                    let mut path: Vec<&str> = ref_name_string.split('/').collect();
                    if let Some(branch_short) = path.pop() {
                        let s = branch_short.to_string();
                        log!(self, "Last part of full name is: {}", s);
                        return Some(s);
                    }
                }
            }
        }

        // we're getting the commit hash as name
        match r.target() {
            Some(v) => {
                let hash_str = v.to_string();
                if hash_str.len() >= 8 {
                    let (s, _) = hash_str.split_at(7);
                    Some(s.to_string())
                } else {
                    Some(hash_str)
                }
            },
            None => None,
        }
    }

    pub fn get_current_branch_name(&mut self) -> Option<String> {
        if self.cache.is_current_branch_set() {
            return self.cache.get_current_branch();
        }
        let mut head: Option<Reference> = None;
        if head.is_none() {
            head = self.get_head();
        }
        let h = match self.resolve_symbolic_reference(head) {
            Some(v) => v,
            None => return None,
        };
        let current_branch_name = self.get_branch_name_for_reference(h);
        self.cache.set_current_branch_name(current_branch_name.clone());
        current_branch_name
    }

    fn get_branch_remote(&self, reference: Reference) -> Option<RefPair> {
        let b = Branch::wrap(reference);
        let upstream = match b.upstream() {
            Ok(u) => u,
            Err(e) => {
                log!(self, "Can't get upstream branch for {:?}", b.name());
                return None;
            }
        };
        let branch_name = match upstream.name() {
            Ok(o_n) => match o_n {
                Some(n) => n,
                None => {
                    log!(self, "Invalid branch name");
                    return None;
                },
            },
            Err(e) => {
                log!(self, "Error while getting name for upstream branch");
                return None;
            }
        };
        let upstream_reference = upstream.get();
        let oid = match upstream_reference.target() {
            Some(o) => o,
            None => {
                log!(self, "Can't get oid of upstream branch");
                return None;
            }
        };
        let remote_name = match upstream_reference.name() {
            Some(n) => {
                let v: Vec<&str> = n.splitn(3, "/").collect();
                if v.len() >= 3 {
                    v[1]
                } else {
                    log!(self, "Can't figure out remote name: {:?}", v);
                    return None;
                }
            },
            None => {
                log!(self, "Can't get full name of upstream branch.");
                return None;
            }
        };
        Some(RefPair{ remote_name: remote_name.to_string(),
                      branch_name: branch_name.to_string(), oid: oid })
    }

    fn get_current_branch_remote_oid(&self) -> Option<RefPair> {
        match self.get_head() {
            Some(r) => self.get_branch_remote(r),
            None => None
        }
    }

    pub fn get_branch_ahead_behind(&mut self, remote_branch: Option<RemoteBranch>) -> Option<BranchAheadBehind> {
        let current_branch_name = self.get_current_branch_name();
        log!(self, "Current branch name = {:?}", current_branch_name);
        let mut ab = BranchAheadBehind::new(current_branch_name);
        let ref_pair_option = self.get_remote_branch(remote_branch);
        let ref_pair = match ref_pair_option {
            Some(u) => {
                u.clone()
            },
            None => return Some(ab),
        };
        ab.remote_branch_name = Some(ref_pair.branch_name.clone());
        ab.remote_name = Some(ref_pair.remote_name);

        let oid = match self.get_current_branch_oid() {
            Some(r) => r,
            None => return None
        };
        let res = self.repo.graph_ahead_behind(oid, ref_pair.oid);
        match res {
            Ok((a, b)) => {
                ab.ahead = a;
                ab.behind = b;
            },
            Err(e) => {
                log!(self, "Can't get ahead & behind stats for branch {}", ref_pair.branch_name);
            }
        };
        Some(ab)
    }

    // find remote branch if branch_name is specified
    // if not, get remote tracking branch for current branch
    fn get_remote_branch(&self, remote_branch: Option<RemoteBranch>) -> Option<RefPair> {
        match remote_branch {
            Some(b) => {
                match self.repo.find_branch(&b.remote_branch, BranchType::Remote) {
                    Ok(o) => Some(RefPair{ branch_name: b.remote_branch_name,
                                           remote_name: b.remote_name,
                                           oid: o.into_reference().target().unwrap() }),
                    Err(e) => {
                        // don't panic here - it doesn't exist, we don't care
                        log!(self, "No remote branch found for {}", b.remote_branch_name);
                        None
                    }
                }
            },
            None => {
                self.get_current_branch_remote_oid()
            }
        }
    }

    pub fn get_status(&self) -> Option<Statuses> {
        let mut so = StatusOptions::new();
        let mut opts = so.show(StatusShow::IndexAndWorkdir);
        opts.include_untracked(true);
        match self.repo.statuses(Some(&mut opts)) {
            Ok(s) => Some(s),
            Err(e) => {
                log!(self, "Unable to get status of repository: {:?}", e);
                None
            }
        }
    }

    pub fn get_repository_state(&self) -> String {
        let state = self.repo.state();
        match state {
            RepositoryState::Clean => String::from(""),  // XXX: this seems to be hardcoded
            RepositoryState::Merge => String::from("merge"),
            RepositoryState::Revert | RepositoryState::RevertSequence => String::from("revert"),
            RepositoryState::CherryPick | RepositoryState::CherryPickSequence => String::from("cherry-pick"),
            RepositoryState::Bisect => String::from("bisect"),
            RepositoryState::Rebase |
                RepositoryState::RebaseInteractive |
                RepositoryState::RebaseMerge => String::from("rebase"),
            RepositoryState::ApplyMailbox | RepositoryState::ApplyMailboxOrRebase => String::from("apply"),
        }
    }

    pub fn get_file_status(&self) -> Option<HashMap<&str, u32>> {
        let mut d = HashMap::new();

        let changed = STATUS_WT_MODIFIED | STATUS_WT_DELETED | STATUS_WT_TYPECHANGE | STATUS_WT_RENAMED;
        let staged = STATUS_INDEX_MODIFIED | STATUS_INDEX_DELETED | STATUS_INDEX_TYPECHANGE | STATUS_INDEX_RENAMED | STATUS_INDEX_NEW;

        let statuses = match self.get_status() {
            Some(x) => x,
            None => return None,
        };

        for s in statuses.iter() {
            let file_status = s.status();
            log!(self, "{}", s.path().unwrap());

            if file_status.intersects(changed) {
                let counter = d.entry(CHANGED_KEY).or_insert(0);
                *counter += 1;
            };
            if file_status.contains(STATUS_WT_NEW) {
                let counter = d.entry(NEW_KEY).or_insert(0);
                *counter += 1;
            };
            if file_status.intersects(staged) {
                let counter = d.entry(STAGED_KEY).or_insert(0);
                *counter += 1;
            };
            if file_status.intersects(STATUS_CONFLICTED) {
                let counter = d.entry(CONFLICTS_KEY).or_insert(0);
                *counter += 1;
            };
        }
        Some(d)
    }
}
<commit_msg>clippy: unused vars<commit_after>use std::collections::HashMap;
use std::cell::RefCell;

use conf::{RemoteBranch};
use constants::{CHANGED_KEY,NEW_KEY,STAGED_KEY,CONFLICTS_KEY};

use git2::*;

struct Cache {
    current_branch_name: RefCell<Option<String>>,
    // TODO: Reference can't be cached (can't be cloned)
    //       implement via OID
    // head: RefCell<Option<Oid>>,
}

pub struct Backend {
    cache: Cache,
    pub repo: Repository,
    pub debug: bool,
}

pub struct BranchAheadBehind {
    pub local_branch_name: Option<String>,
    pub remote_branch_name: Option<String>,
    pub remote_name: Option<String>,
    pub ahead: usize,
    pub behind: usize
}

impl BranchAheadBehind {
    fn new(l: Option<String>) -> BranchAheadBehind {
        BranchAheadBehind{ local_branch_name: l, remote_branch_name: None, remote_name: None,
                           ahead: 0, behind: 0 }
    }
}

#[derive(Clone)]
struct RefPair {
    remote_name: String,
    branch_name: String,
    oid: Oid,
}

impl Cache {
    pub fn new() -> Cache {
        Cache{ current_branch_name: RefCell::new(None) }
    }

    fn set_current_branch_name(&self, n: Option<String>) {
        let mut c = self.current_branch_name.borrow_mut();
        *c = n;
    }

    fn is_current_branch_set(&self) -> bool {
        self.current_branch_name.borrow().is_some()
    }

    fn get_current_branch(&self) -> Option<String> {
        self.current_branch_name.borrow().clone()
    }
}

impl Backend {
    pub fn new(repo: Repository, debug: bool) -> Backend {
        Backend{ repo: repo, debug: debug, cache: Cache::new() }
    }

    fn get_head(&self) -> Option<Reference> {
        match self.repo.head() {
            Ok(head) => Some(head),
            Err(e2) => {
                log!(self, "Can't get HEAD: {}", e2);
                match self.repo.find_reference("HEAD") {
                    Ok(x) => {
                        log!(self, "Found HEAD directly: {:?}", x.name());
                        Some(x)
                    },
                    Err(e) => {
                        log!(self, "reference HEAD: {}", e);
                        return None
                    },
                }
            },
        }
    }

    fn get_current_branch_oid(&self) -> Option<Oid> {
        let head = self.get_head();
        match head {
            Some(v) => v.target(),
            None => {
                log!(self, "Failed to find Oid for HEAD.");
                return None
            },
        }
    }

    fn resolve_symbolic_reference<'a>(&self, reference: Option<Reference<'a>>) -> Option<Reference<'a>> {
        let s_reference = match reference {
            Some(v) => {
                match v.resolve() {
                    Ok(y) => {
                        log!(self, "Branch name is {:?}", y.name());
                        y
                    },
                    Err(e) => {
                        log!(self, "Branch name is {:?}, error: {:?}", v.name(), e);
                        v
                    }
                }
            }
            None => {
                log!(self, "No branch name found");
                return None;
            }
        };
        Some(s_reference)
    }

    fn get_branch_name_for_reference(&self, r: Reference) -> Option<String> {
        if let Some(ref_name_string) = r.shorthand() {
            if ref_name_string != "HEAD" {
                let s = ref_name_string.to_string();
                log!(self, "Shorthand for reference is: {}", s);
                return Some(s);
            } else {
                if let Some(ref_name) = r.symbolic_target() {
                    let ref_name_string = ref_name.to_string();
                    log!(self, "shorthand = HEAD, links to: {}", ref_name_string);
                    let mut path: Vec<&str> = ref_name_string.split('/').collect();
                    if let Some(branch_short) = path.pop() {
                        let s = branch_short.to_string();
                        log!(self, "Last part of full name is: {}", s);
                        return Some(s);
                    }
                }
            }
        }

        // we're getting the commit hash as name
        match r.target() {
            Some(v) => {
                let hash_str = v.to_string();
                if hash_str.len() >= 8 {
                    let (s, _) = hash_str.split_at(7);
                    Some(s.to_string())
                } else {
                    Some(hash_str)
                }
            },
            None => None,
        }
    }

    pub fn get_current_branch_name(&mut self) -> Option<String> {
        if self.cache.is_current_branch_set() {
            return self.cache.get_current_branch();
        }
        let mut head: Option<Reference> = None;
        if head.is_none() {
            head = self.get_head();
        }
        let h = match self.resolve_symbolic_reference(head) {
            Some(v) => v,
            None => return None,
        };
        let current_branch_name = self.get_branch_name_for_reference(h);
        self.cache.set_current_branch_name(current_branch_name.clone());
        current_branch_name
    }

    fn get_branch_remote(&self, reference: Reference) -> Option<RefPair> {
        let b = Branch::wrap(reference);
        let upstream = match b.upstream() {
            Ok(u) => u,
            Err(e) => {
                log!(self, "Can't get upstream branch for {:?}: {:?}", b.name(), e);
                return None;
            }
        };
        let branch_name = match upstream.name() {
            Ok(o_n) => match o_n {
                Some(n) => n,
                None => {
                    log!(self, "Invalid branch name");
                    return None;
                },
            },
            Err(e) => {
                log!(self, "Error while getting name for upstream branch: {:?}", e);
                return None;
            }
        };
        let upstream_reference = upstream.get();
        let oid = match upstream_reference.target() {
            Some(o) => o,
            None => {
                log!(self, "Can't get oid of upstream branch");
                return None;
            }
        };
        let remote_name = match upstream_reference.name() {
            Some(n) => {
                let v: Vec<&str> = n.splitn(3, "/").collect();
                if v.len() >= 3 {
                    v[1]
                } else {
                    log!(self, "Can't figure out remote name: {:?}", v);
                    return None;
                }
            },
            None => {
                log!(self, "Can't get full name of upstream branch.");
                return None;
            }
        };
        Some(RefPair{ remote_name: remote_name.to_string(),
                      branch_name: branch_name.to_string(), oid: oid })
    }

    fn get_current_branch_remote_oid(&self) -> Option<RefPair> {
        match self.get_head() {
            Some(r) => self.get_branch_remote(r),
            None => None
        }
    }

    pub fn get_branch_ahead_behind(&mut self, remote_branch: Option<RemoteBranch>) -> Option<BranchAheadBehind> {
        let current_branch_name = self.get_current_branch_name();
        log!(self, "Current branch name = {:?}", current_branch_name);
        let mut ab = BranchAheadBehind::new(current_branch_name);
        let ref_pair_option = self.get_remote_branch(remote_branch);
        let ref_pair = match ref_pair_option {
            Some(u) => {
                u.clone()
            },
            None => return Some(ab),
        };
        ab.remote_branch_name = Some(ref_pair.branch_name.clone());
        ab.remote_name = Some(ref_pair.remote_name);

        let oid = match self.get_current_branch_oid() {
            Some(r) => r,
            None => return None
        };
        let res = self.repo.graph_ahead_behind(oid, ref_pair.oid);
        match res {
            Ok((a, b)) => {
                ab.ahead = a;
                ab.behind = b;
            },
            Err(e) => {
                log!(self, "Can't get ahead & behind stats for branch {}: {:?}", ref_pair.branch_name, e);
            }
        };
        Some(ab)
    }

    // find remote branch if branch_name is specified
    // if not, get remote tracking branch for current branch
    fn get_remote_branch(&self, remote_branch: Option<RemoteBranch>) -> Option<RefPair> {
        match remote_branch {
            Some(b) => {
                match self.repo.find_branch(&b.remote_branch, BranchType::Remote) {
                    Ok(o) => Some(RefPair{ branch_name: b.remote_branch_name,
                                           remote_name: b.remote_name,
                                           oid: o.into_reference().target().unwrap() }),
                    Err(e) => {
                        // don't panic here - it doesn't exist, we don't care
                        log!(self, "No remote branch found for {}: {:?}", b.remote_branch_name, e);
                        None
                    }
                }
            },
            None => {
                self.get_current_branch_remote_oid()
            }
        }
    }

    pub fn get_status(&self) -> Option<Statuses> {
        let mut so = StatusOptions::new();
        let mut opts = so.show(StatusShow::IndexAndWorkdir);
        opts.include_untracked(true);
        match self.repo.statuses(Some(&mut opts)) {
            Ok(s) => Some(s),
            Err(e) => {
                log!(self, "Unable to get status of repository: {:?}", e);
                None
            }
        }
    }

    pub fn get_repository_state(&self) -> String {
        let state = self.repo.state();
        match state {
            RepositoryState::Clean => String::from(""),  // XXX: this seems to be hardcoded
            RepositoryState::Merge => String::from("merge"),
            RepositoryState::Revert | RepositoryState::RevertSequence => String::from("revert"),
            RepositoryState::CherryPick | RepositoryState::CherryPickSequence => String::from("cherry-pick"),
            RepositoryState::Bisect => String::from("bisect"),
            RepositoryState::Rebase |
                RepositoryState::RebaseInteractive |
                RepositoryState::RebaseMerge => String::from("rebase"),
            RepositoryState::ApplyMailbox | RepositoryState::ApplyMailboxOrRebase => String::from("apply"),
        }
    }

    pub fn get_file_status(&self) -> Option<HashMap<&str, u32>> {
        let mut d = HashMap::new();

        let changed = STATUS_WT_MODIFIED | STATUS_WT_DELETED | STATUS_WT_TYPECHANGE | STATUS_WT_RENAMED;
        let staged = STATUS_INDEX_MODIFIED | STATUS_INDEX_DELETED | STATUS_INDEX_TYPECHANGE | STATUS_INDEX_RENAMED | STATUS_INDEX_NEW;

        let statuses = match self.get_status() {
            Some(x) => x,
            None => return None,
        };

        for s in statuses.iter() {
            let file_status = s.status();
            log!(self, "{}", s.path().unwrap());

            if file_status.intersects(changed) {
                let counter = d.entry(CHANGED_KEY).or_insert(0);
                *counter += 1;
            };
            if file_status.contains(STATUS_WT_NEW) {
                let counter = d.entry(NEW_KEY).or_insert(0);
                *counter += 1;
            };
            if file_status.intersects(staged) {
                let counter = d.entry(STAGED_KEY).or_insert(0);
                *counter += 1;
            };
            if file_status.intersects(STATUS_CONFLICTED) {
                let counter = d.entry(CONFLICTS_KEY).or_insert(0);
                *counter += 1;
            };
        }
        Some(d)
    }
}
<|endoftext|>


<commit_before>#![allow(unused_imports, unused_must_use, unused_variables, dead_code)]

extern crate combine;

use std::io::prelude::*;
use std::fs::File;
use std::path::Path;
use std::str;
use std::marker::PhantomData;
use combine::{spaces, parser, between, many, many1, digit, char, any, string, token, Parser, ParserExt, ParseError};
use combine::primitives::{State, Stream, ParseResult, Consumed};
use combine::combinator::{Between, Token, FnParser, satisfy};

//my own bencode stuff!
//
pub enum Bencode {
    Number(i64),
    String(String)
}

fn open_file <P: AsRef<Path>>(path: P) -> Vec<u8> {
    let mut fd = File::open(path).unwrap();
    let mut buffer:Vec<u8> = Vec::new();
    fd.read_to_end(&mut buffer);
    buffer
}

enum Expr {
    Integer(i64),
    String(String)
}

//TODO: this does not handle negatives, yo
fn bencode_integer<I>(input: State<I>) -> ParseResult<i64, I> where I:Stream<Item=char> {
    let (open, close) = (char('i'), char('e'));
    let mut int = between(open, close, many1::<String, _>(digit())).map(|x| {
        x.parse::<i64>().unwrap()
    });
    int.parse_state(input)
}

fn bencode_string<I>(input: State<I>) -> ParseResult<String, I> where I:Stream<Item=char> {
    let (len, input_) = try!(bencode_string_length_prefix(input));
    input_.combine(|input__| take(len).parse_state(input__))
}

fn bencode_string_length_prefix<I>(input: State<I>) -> ParseResult<i32, I> where I:Stream<Item=char> {
    let many_digit = many1::<String, _>(digit());
    let mut get_len = (many_digit, char(':')).map(|(length, _)| length.parse::<i32>().unwrap());
    get_len.parse_state(input)
}

//for now only handle i64s
fn bencode_list<I>(input: State<I>) -> ParseResult<Vec<i64>, I> where I:Stream<Item=char> {
    let (open, close) = (char('l'), char('e'));
    let list_contents = many(parser(bencode_integer));
    let mut list = between(open, close, list_contents).map(|x| {
        x
    });
    list.parse_state(input)
}

fn take <I> (num: i32) -> SizedBuffer<I> where I: Stream<Item=char> {
    SizedBuffer(num, PhantomData)
}

#[derive(Clone)]
pub struct SizedBuffer <I>(i32, PhantomData<I>);
impl <I> Parser for SizedBuffer<I> where I: Stream<Item=char> {
    type Input = I;
    type Output = String;
    fn parse_lazy(&mut self, mut input: State<I>) -> ParseResult<String, I> {
        let start = input.position;
        let mut vec_buf:Vec<char> = Vec::new();
        for _ in 0..self.0 {
            match input.uncons() {
                Ok((other, rest)) => {
                    vec_buf.push(other);
                    input = rest.into_inner();
                }
                Err(error) => {
                    return error.combine(|mut error| {
                        error.position = start;
                        Err(Consumed::Consumed(error))
                    })
                }
            };
        }
        Ok((vec_buf.into_iter().collect(), Consumed::Consumed(input)))
    }
}

fn main () {
    let file_contents = open_file("t-test.torrent");

    match str::from_utf8(&file_contents) {
        Ok(v) => println!("file: {}", v),
        Err(e) => panic!("Not a UTF-8 String: {}", e)
    }
}

#[test]
fn test_integer() {
    let result = parser(bencode_integer).parse("i57e");
    assert_eq!(result, Ok((57,"")));
}

#[test]
fn test_string() {
    let result = parser(bencode_string).parse("5:abcde");
    assert_eq!(result, Ok(("abcde".to_string(), "")));;
}

//#[test]
fn test_list_integer() {
    let result = parser(bencode_list).parse("li57ei32ee");
    assert_eq!(result, Ok((vec![57, 32], "")));
}
<commit_msg>got bencode list to work with homogenous, heterogenous int/string sequences<commit_after>#![allow(unused_imports, unused_must_use, unused_variables, dead_code)]

extern crate combine;

use std::io::prelude::*;
use std::fs::File;
use std::path::Path;
use std::str;
use std::marker::PhantomData;
use combine::{spaces, parser, between, many, many1, digit, char, any, string, token, Parser, ParserExt, ParseError};
use combine::primitives::{State, Stream, ParseResult, Consumed};
use combine::combinator::{Between, Token, FnParser, satisfy};

//my own bencode stuff!
//
fn open_file <P: AsRef<Path>>(path: P) -> Vec<u8> {
    let mut fd = File::open(path).unwrap();
    let mut buffer:Vec<u8> = Vec::new();
    fd.read_to_end(&mut buffer);
    buffer
}

#[derive(Debug, Eq, PartialEq)]
enum Bencode {
    Int(i64),
    String(String)
}

//TODO: this does not handle negatives, yo
fn bencode_integer<I>(input: State<I>) -> ParseResult<Bencode, I> where I:Stream<Item=char> {
    let (open, close) = (char('i'), char('e'));
    let mut int = between(open, close, many1::<String, _>(digit())).map(|x| {
        Bencode::Int(x.parse::<i64>().unwrap())
    });
    int.parse_state(input)
}

fn bencode_string<I>(input: State<I>) -> ParseResult<Bencode, I> where I:Stream<Item=char> {
    let (len, input_) = try!(bencode_string_length_prefix(input));
    input_.combine(|input__| take(len).map(|x| Bencode::String(x)).parse_state(input__))
}

fn bencode_string_length_prefix<I>(input: State<I>) -> ParseResult<i32, I> where I:Stream<Item=char> {
    let many_digit = many1::<String, _>(digit());
    let mut get_len = (many_digit, char(':')).map(|(length, _)| length.parse::<i32>().unwrap());
    get_len.parse_state(input)
}

//for now only handle i64s
fn bencode_list<I>(input: State<I>) -> ParseResult<Vec<Bencode>, I> where I:Stream<Item=char> {
    let (open, close) = (char('l'), char('e'));
    let list_contents = many(parser(bencode_integer).or(parser(bencode_string)));
    let mut list = between(open, close, list_contents).map(|x| {
        x
    });
    list.parse_state(input)
}

fn take <I> (num: i32) -> SizedBuffer<I> where I: Stream<Item=char> {
    SizedBuffer(num, PhantomData)
}

#[derive(Clone)]
pub struct SizedBuffer <I>(i32, PhantomData<I>);
impl <I> Parser for SizedBuffer<I> where I: Stream<Item=char> {
    type Input = I;
    type Output = String;
    fn parse_lazy(&mut self, mut input: State<I>) -> ParseResult<String, I> {
        let start = input.position;
        let mut vec_buf:Vec<char> = Vec::new();
        for _ in 0..self.0 {
            match input.uncons() {
                Ok((other, rest)) => {
                    vec_buf.push(other);
                    input = rest.into_inner();
                }
                Err(error) => {
                    return error.combine(|mut error| {
                        error.position = start;
                        Err(Consumed::Consumed(error))
                    })
                }
            };
        }
        Ok((vec_buf.into_iter().collect(), Consumed::Consumed(input)))
    }
}

fn main () {
    let file_contents = open_file("t-test.torrent");

    match str::from_utf8(&file_contents) {
        Ok(v) => println!("file: {}", v),
        Err(e) => panic!("Not a UTF-8 String: {}", e)
    }
}

#[test]
fn test_integer() {
    let result = parser(bencode_integer).parse("i57e");
    assert_eq!(result, Ok((Bencode::Int(57),"")));
}

#[test]
fn test_string() {
    let result = parser(bencode_string).parse("5:abcde");
    assert_eq!(result, Ok((Bencode::String("abcde".to_string()), "")));;
}

#[test]
fn test_list() {
    let homogenous_int_list = parser(bencode_list).parse("li57ei32ee");
    assert_eq!(homogenous_int_list, Ok((vec![Bencode::Int(57), Bencode::Int(32)], "")));

    let homogenous_str_list = parser(bencode_list).parse("l3:abc5:defghe");
    assert_eq!(homogenous_str_list, Ok((vec![Bencode::String("abc".to_string()), Bencode::String("defgh".to_string())], "")));

    let hetero_list = parser(bencode_list).parse("li32e3:abce");
    assert_eq!(hetero_list, Ok((vec![Bencode::Int(32), Bencode::String("abc".to_string())], "")));
}

<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

//! An implementation of re-entrant mutexes.
//!
//! Re-entrant mutexes are like mutexes, but where it is expected
//! that a single thread may own a lock more than once.

//! It provides the same interface as https://github.com/rust-lang/rust/blob/master/src/libstd/sys/common/remutex.rs
//! so if those types are ever exported, we should be able to replace this implemtation.

#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;

use std::cell::{Cell, UnsafeCell};
use std::num::NonZeroUsize;
use std::ops::Deref;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{LockResult, Mutex, MutexGuard, PoisonError, TryLockError, TryLockResult};

/// A type for thread ids.

// TODO: can we use the thread-id crate for this?

#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct ThreadId(NonZeroUsize);

lazy_static! {
    static ref THREAD_COUNT: AtomicUsize = AtomicUsize::new(1);
}

impl ThreadId {
    #[allow(unsafe_code)]
    fn new() -> ThreadId {
        let number = THREAD_COUNT.fetch_add(1, Ordering::SeqCst);
        ThreadId(NonZeroUsize::new(number).unwrap())
    }
    pub fn current() -> ThreadId {
        THREAD_ID.with(|tls| tls.clone())
    }
}

thread_local! { static THREAD_ID: ThreadId = ThreadId::new() }

/// A type for atomic storage of thread ids.
#[derive(Debug)]
pub struct AtomicOptThreadId(AtomicUsize);

impl AtomicOptThreadId {
    pub fn new() -> AtomicOptThreadId {
        AtomicOptThreadId(AtomicUsize::new(0))
    }
    pub fn store(&self, value: Option<ThreadId>, ordering: Ordering) {
        let number = value.map(|id| id.0.get()).unwrap_or(0);
        self.0.store(number, ordering);
    }
    #[allow(unsafe_code)]
    pub fn load(&self, ordering: Ordering) -> Option<ThreadId> {
        let number = self.0.load(ordering);
        NonZeroUsize::new(number).map(ThreadId)
    }
}

/// A type for hand-over-hand mutexes.
///
/// These support `lock` and `unlock` functions. `lock` blocks waiting to become the
/// mutex owner. `unlock` can only be called by the lock owner, and panics otherwise.
/// They have the same happens-before and poisoning semantics as `Mutex`.

// TODO: Can we use `raw_lock` and `raw_unlock` from `parking_lot`'s `Mutex` for this?

pub struct HandOverHandMutex {
    mutex: Mutex<()>,
    owner: AtomicOptThreadId,
    guard: UnsafeCell<Option<MutexGuard<'static, ()>>>,
}

impl HandOverHandMutex {
    pub fn new() -> HandOverHandMutex {
        HandOverHandMutex {
            mutex: Mutex::new(()),
            owner: AtomicOptThreadId::new(),
            guard: UnsafeCell::new(None),
        }
    }
    #[allow(unsafe_code)]
    unsafe fn set_guard_and_owner<'a>(&'a self, guard: MutexGuard<'a, ()>) {
        // The following two lines allow us to unsafely store
        // Some(guard): Option<MutexGuard<'a, ()>
        // in self.guard, even though its contents are Option<MutexGuard<'static, ()>>,
        // that is the lifetime is 'a not 'static.
        let guard_ptr = &mut *(self.guard.get() as *mut u8 as *mut Option<MutexGuard<'a, ()>>);
        *guard_ptr = Some(guard);
        self.owner
            .store(Some(ThreadId::current()), Ordering::Relaxed);
    }
    #[allow(unsafe_code)]
    unsafe fn unset_guard_and_owner(&self) {
        let guard_ptr = &mut *self.guard.get();
        let old_owner = self.owner();
        self.owner.store(None, Ordering::Relaxed);
        // Make sure we release the lock before checking the assertions.
        // We protect logging by a re-entrant lock, so we don't want
        // to do any incidental logging while we the lock is held.
        drop(guard_ptr.take());
        // Now we have released the lock, it's okay to use logging.
        assert_eq!(old_owner, Some(ThreadId::current()));
    }
    #[allow(unsafe_code)]
    pub fn lock<'a>(&'a self) -> LockResult<()> {
        let (guard, result) = match self.mutex.lock() {
            Ok(guard) => (guard, Ok(())),
            Err(err) => (err.into_inner(), Err(PoisonError::new(()))),
        };
        unsafe {
            self.set_guard_and_owner(guard);
        }
        result
    }
    #[allow(unsafe_code)]
    pub fn try_lock(&self) -> TryLockResult<()> {
        let (guard, result) = match self.mutex.try_lock() {
            Ok(guard) => (guard, Ok(())),
            Err(TryLockError::WouldBlock) => return Err(TryLockError::WouldBlock),
            Err(TryLockError::Poisoned(err)) => (
                err.into_inner(),
                Err(TryLockError::Poisoned(PoisonError::new(()))),
            ),
        };
        unsafe {
            self.set_guard_and_owner(guard);
        }
        result
    }
    #[allow(unsafe_code)]
    pub fn unlock(&self) {
        unsafe {
            self.unset_guard_and_owner();
        }
    }
    pub fn owner(&self) -> Option<ThreadId> {
        self.owner.load(Ordering::Relaxed)
    }
}

#[allow(unsafe_code)]
unsafe impl Send for HandOverHandMutex {}

/// A type for re-entrant mutexes.
///
/// It provides the same interface as https://github.com/rust-lang/rust/blob/master/src/libstd/sys/common/remutex.rs

pub struct ReentrantMutex<T> {
    mutex: HandOverHandMutex,
    count: Cell<usize>,
    data: T,
}

#[allow(unsafe_code)]
unsafe impl<T> Sync for ReentrantMutex<T> where T: Send {}

impl<T> ReentrantMutex<T> {
    pub fn new(data: T) -> ReentrantMutex<T> {
        trace!("{:?} Creating new lock.", ThreadId::current());
        ReentrantMutex {
            mutex: HandOverHandMutex::new(),
            count: Cell::new(0),
            data: data,
        }
    }

    pub fn lock(&self) -> LockResult<ReentrantMutexGuard<T>> {
        trace!("{:?} Locking.", ThreadId::current());
        if self.mutex.owner() != Some(ThreadId::current()) {
            trace!("{:?} Becoming owner.", ThreadId::current());
            if let Err(_) = self.mutex.lock() {
                trace!("{:?} Poison!", ThreadId::current());
                return Err(PoisonError::new(self.mk_guard()));
            }
            trace!("{:?} Became owner.", ThreadId::current());
        }
        Ok(self.mk_guard())
    }

    pub fn try_lock(&self) -> TryLockResult<ReentrantMutexGuard<T>> {
        trace!("{:?} Try locking.", ThreadId::current());
        if self.mutex.owner() != Some(ThreadId::current()) {
            trace!("{:?} Becoming owner?", ThreadId::current());
            if let Err(err) = self.mutex.try_lock() {
                match err {
                    TryLockError::WouldBlock => {
                        trace!("{:?} Would block.", ThreadId::current());
                        return Err(TryLockError::WouldBlock);
                    },
                    TryLockError::Poisoned(_) => {
                        trace!("{:?} Poison!", ThreadId::current());
                        return Err(TryLockError::Poisoned(PoisonError::new(self.mk_guard())));
                    },
                }
            }
            trace!("{:?} Became owner.", ThreadId::current());
        }
        Ok(self.mk_guard())
    }

    fn unlock(&self) {
        trace!("{:?} Unlocking.", ThreadId::current());
        let count = self
            .count
            .get()
            .checked_sub(1)
            .expect("Underflowed lock count.");
        trace!("{:?} Decrementing count to {}.", ThreadId::current(), count);
        self.count.set(count);
        if count == 0 {
            trace!("{:?} Releasing mutex.", ThreadId::current());
            self.mutex.unlock();
        }
    }

    fn mk_guard(&self) -> ReentrantMutexGuard<T> {
        let count = self
            .count
            .get()
            .checked_add(1)
            .expect("Overflowed lock count.");
        trace!("{:?} Incrementing count to {}.", ThreadId::current(), count);
        self.count.set(count);
        ReentrantMutexGuard { mutex: self }
    }
}

#[must_use]
pub struct ReentrantMutexGuard<'a, T>
where
    T: 'static,
{
    mutex: &'a ReentrantMutex<T>,
}

impl<'a, T> Drop for ReentrantMutexGuard<'a, T> {
    #[allow(unsafe_code)]
    fn drop(&mut self) {
        self.mutex.unlock()
    }
}

impl<'a, T> Deref for ReentrantMutexGuard<'a, T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.mutex.data
    }
}
<commit_msg>Fix rust's remutex source code link<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

//! An implementation of re-entrant mutexes.
//!
//! Re-entrant mutexes are like mutexes, but where it is expected
//! that a single thread may own a lock more than once.

//! It provides the same interface as https://github.com/rust-lang/rust/blob/5edaa7eefd76d4996dcf85dfc1c1a3f737087257/src/libstd/sys_common/remutex.rs
//! so if those types are ever exported, we should be able to replace this implemtation.

#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;

use std::cell::{Cell, UnsafeCell};
use std::num::NonZeroUsize;
use std::ops::Deref;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{LockResult, Mutex, MutexGuard, PoisonError, TryLockError, TryLockResult};

/// A type for thread ids.

// TODO: can we use the thread-id crate for this?

#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct ThreadId(NonZeroUsize);

lazy_static! {
    static ref THREAD_COUNT: AtomicUsize = AtomicUsize::new(1);
}

impl ThreadId {
    #[allow(unsafe_code)]
    fn new() -> ThreadId {
        let number = THREAD_COUNT.fetch_add(1, Ordering::SeqCst);
        ThreadId(NonZeroUsize::new(number).unwrap())
    }
    pub fn current() -> ThreadId {
        THREAD_ID.with(|tls| tls.clone())
    }
}

thread_local! { static THREAD_ID: ThreadId = ThreadId::new() }

/// A type for atomic storage of thread ids.
#[derive(Debug)]
pub struct AtomicOptThreadId(AtomicUsize);

impl AtomicOptThreadId {
    pub fn new() -> AtomicOptThreadId {
        AtomicOptThreadId(AtomicUsize::new(0))
    }
    pub fn store(&self, value: Option<ThreadId>, ordering: Ordering) {
        let number = value.map(|id| id.0.get()).unwrap_or(0);
        self.0.store(number, ordering);
    }
    #[allow(unsafe_code)]
    pub fn load(&self, ordering: Ordering) -> Option<ThreadId> {
        let number = self.0.load(ordering);
        NonZeroUsize::new(number).map(ThreadId)
    }
}

/// A type for hand-over-hand mutexes.
///
/// These support `lock` and `unlock` functions. `lock` blocks waiting to become the
/// mutex owner. `unlock` can only be called by the lock owner, and panics otherwise.
/// They have the same happens-before and poisoning semantics as `Mutex`.

// TODO: Can we use `raw_lock` and `raw_unlock` from `parking_lot`'s `Mutex` for this?

pub struct HandOverHandMutex {
    mutex: Mutex<()>,
    owner: AtomicOptThreadId,
    guard: UnsafeCell<Option<MutexGuard<'static, ()>>>,
}

impl HandOverHandMutex {
    pub fn new() -> HandOverHandMutex {
        HandOverHandMutex {
            mutex: Mutex::new(()),
            owner: AtomicOptThreadId::new(),
            guard: UnsafeCell::new(None),
        }
    }
    #[allow(unsafe_code)]
    unsafe fn set_guard_and_owner<'a>(&'a self, guard: MutexGuard<'a, ()>) {
        // The following two lines allow us to unsafely store
        // Some(guard): Option<MutexGuard<'a, ()>
        // in self.guard, even though its contents are Option<MutexGuard<'static, ()>>,
        // that is the lifetime is 'a not 'static.
        let guard_ptr = &mut *(self.guard.get() as *mut u8 as *mut Option<MutexGuard<'a, ()>>);
        *guard_ptr = Some(guard);
        self.owner
            .store(Some(ThreadId::current()), Ordering::Relaxed);
    }
    #[allow(unsafe_code)]
    unsafe fn unset_guard_and_owner(&self) {
        let guard_ptr = &mut *self.guard.get();
        let old_owner = self.owner();
        self.owner.store(None, Ordering::Relaxed);
        // Make sure we release the lock before checking the assertions.
        // We protect logging by a re-entrant lock, so we don't want
        // to do any incidental logging while we the lock is held.
        drop(guard_ptr.take());
        // Now we have released the lock, it's okay to use logging.
        assert_eq!(old_owner, Some(ThreadId::current()));
    }
    #[allow(unsafe_code)]
    pub fn lock<'a>(&'a self) -> LockResult<()> {
        let (guard, result) = match self.mutex.lock() {
            Ok(guard) => (guard, Ok(())),
            Err(err) => (err.into_inner(), Err(PoisonError::new(()))),
        };
        unsafe {
            self.set_guard_and_owner(guard);
        }
        result
    }
    #[allow(unsafe_code)]
    pub fn try_lock(&self) -> TryLockResult<()> {
        let (guard, result) = match self.mutex.try_lock() {
            Ok(guard) => (guard, Ok(())),
            Err(TryLockError::WouldBlock) => return Err(TryLockError::WouldBlock),
            Err(TryLockError::Poisoned(err)) => (
                err.into_inner(),
                Err(TryLockError::Poisoned(PoisonError::new(()))),
            ),
        };
        unsafe {
            self.set_guard_and_owner(guard);
        }
        result
    }
    #[allow(unsafe_code)]
    pub fn unlock(&self) {
        unsafe {
            self.unset_guard_and_owner();
        }
    }
    pub fn owner(&self) -> Option<ThreadId> {
        self.owner.load(Ordering::Relaxed)
    }
}

#[allow(unsafe_code)]
unsafe impl Send for HandOverHandMutex {}

/// A type for re-entrant mutexes.
///
/// It provides the same interface as https://github.com/rust-lang/rust/blob/5edaa7eefd76d4996dcf85dfc1c1a3f737087257/src/libstd/sys_common/remutex.rs

pub struct ReentrantMutex<T> {
    mutex: HandOverHandMutex,
    count: Cell<usize>,
    data: T,
}

#[allow(unsafe_code)]
unsafe impl<T> Sync for ReentrantMutex<T> where T: Send {}

impl<T> ReentrantMutex<T> {
    pub fn new(data: T) -> ReentrantMutex<T> {
        trace!("{:?} Creating new lock.", ThreadId::current());
        ReentrantMutex {
            mutex: HandOverHandMutex::new(),
            count: Cell::new(0),
            data: data,
        }
    }

    pub fn lock(&self) -> LockResult<ReentrantMutexGuard<T>> {
        trace!("{:?} Locking.", ThreadId::current());
        if self.mutex.owner() != Some(ThreadId::current()) {
            trace!("{:?} Becoming owner.", ThreadId::current());
            if let Err(_) = self.mutex.lock() {
                trace!("{:?} Poison!", ThreadId::current());
                return Err(PoisonError::new(self.mk_guard()));
            }
            trace!("{:?} Became owner.", ThreadId::current());
        }
        Ok(self.mk_guard())
    }

    pub fn try_lock(&self) -> TryLockResult<ReentrantMutexGuard<T>> {
        trace!("{:?} Try locking.", ThreadId::current());
        if self.mutex.owner() != Some(ThreadId::current()) {
            trace!("{:?} Becoming owner?", ThreadId::current());
            if let Err(err) = self.mutex.try_lock() {
                match err {
                    TryLockError::WouldBlock => {
                        trace!("{:?} Would block.", ThreadId::current());
                        return Err(TryLockError::WouldBlock);
                    },
                    TryLockError::Poisoned(_) => {
                        trace!("{:?} Poison!", ThreadId::current());
                        return Err(TryLockError::Poisoned(PoisonError::new(self.mk_guard())));
                    },
                }
            }
            trace!("{:?} Became owner.", ThreadId::current());
        }
        Ok(self.mk_guard())
    }

    fn unlock(&self) {
        trace!("{:?} Unlocking.", ThreadId::current());
        let count = self
            .count
            .get()
            .checked_sub(1)
            .expect("Underflowed lock count.");
        trace!("{:?} Decrementing count to {}.", ThreadId::current(), count);
        self.count.set(count);
        if count == 0 {
            trace!("{:?} Releasing mutex.", ThreadId::current());
            self.mutex.unlock();
        }
    }

    fn mk_guard(&self) -> ReentrantMutexGuard<T> {
        let count = self
            .count
            .get()
            .checked_add(1)
            .expect("Overflowed lock count.");
        trace!("{:?} Incrementing count to {}.", ThreadId::current(), count);
        self.count.set(count);
        ReentrantMutexGuard { mutex: self }
    }
}

#[must_use]
pub struct ReentrantMutexGuard<'a, T>
where
    T: 'static,
{
    mutex: &'a ReentrantMutex<T>,
}

impl<'a, T> Drop for ReentrantMutexGuard<'a, T> {
    #[allow(unsafe_code)]
    fn drop(&mut self) {
        self.mutex.unlock()
    }
}

impl<'a, T> Deref for ReentrantMutexGuard<'a, T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.mutex.data
    }
}
<|endoftext|>