<commit_before>// Adapton uses memoization under the covers, which needs an efficient
// mechanism to search for function pointers and them for equality.
//
// Meanwhile, Rust does not provide Eq and Hash implementations for
// trait Fn.  So, to identify Rust functions as values that we can
// hash and compare, we need to bundle additional static information
// along with the function pointer as use this data as a proxy for the
// function itself.  The idea is that this information uniquely
// identifies the function pointer (i.e., two distinct functions will
// always have two distinct identities).
//
// TODO: make a macro for wrapping Fn's as FnObj's, defined below:

use std::fmt::{Debug,Formatter,Result};
use std::hash::{Hash,Hasher};
use std::rc::Rc;
use adapton_sigs::Adapton;

#[derive(Hash,Debug)]
pub struct ProgPt {
    file:Rc<String>,
    module:Rc<String>,
    name:Rc<String>,
    line:usize,
    column:usize,
}

pub struct FnObj<A:Adapton,Arg,Res> {
    prog_pt:ProgPt,
    fn_val:Box<Fn(&mut A, Rc<Arg>) -> Rc<Res>>,
}

impl<A:Adapton,Arg,Res> Debug for FnObj<A,Arg,Res> {
    fn fmt(&self, f: &mut Formatter) -> Result { self.prog_pt.fmt(f) }
}

impl<A:Adapton,Arg,Res> Hash for FnObj<A,Arg,Res> {
    fn hash<H>(&self, state: &mut H) where H: Hasher { self.prog_pt.hash(state) }
    fn hash_slice<H>(data: &[Self], state: &mut H) where H: Hasher { panic!("TODO") }
}
<commit_msg>plans for macro usage, fast comparisons<commit_after>// Adapton uses memoization under the covers, which needs an efficient
// mechanism to search for function pointers and them for equality.
//
// Meanwhile, Rust does not provide Eq and Hash implementations for
// trait Fn.  So, to identify Rust functions as values that we can
// hash and compare, we need to bundle additional static information
// along with the function pointer as use this data as a proxy for the
// function itself.  The idea is that this information uniquely
// identifies the function pointer (i.e., two distinct functions will
// always have two distinct identities).
//
// TODO: make a macro for wrapping Fn's as FnObj's, defined below:

use std::fmt::{Debug,Formatter,Result};
use std::hash::{Hash,Hasher};
use std::rc::Rc;
use adapton_sigs::Adapton;

#[derive(Hash,Debug)]
pub struct ProgPt {
    hash:u64, // hash of all fields below:

    // Symbolic identity, in Rust semantics:
    symbol:Rc<String>, // via stringify!(...)
    module:Rc<String>, // via module!()

    // Location in local filesystem:
    file:Rc<String>,   // via file!()
    line:usize,        // via line!()
    column:usize,      // via column!()
}

pub struct FnObj<A:Adapton,Arg,Res> {
    prog_pt:ProgPt,
    fn_val:Box<Fn(&mut A, Rc<Arg>) -> Rc<Res>>,
}

impl<A:Adapton,Arg,Res> Debug for FnObj<A,Arg,Res> {
    fn fmt(&self, f: &mut Formatter) -> Result { self.prog_pt.fmt(f) }
}

impl<A:Adapton,Arg,Res> Hash for FnObj<A,Arg,Res> {
    fn hash<H>(&self, state: &mut H) where H: Hasher { self.prog_pt.hash(state) }
    fn hash_slice<H>(data: &[Self], state: &mut H) where H: Hasher { panic!("TODO") }
}
<|endoftext|>


<commit_before>#![allow(clippy::too_many_arguments)]

extern crate pq_sys;

use self::pq_sys::*;
use std::ffi::{CStr, CString};
use std::os::raw as libc;
use std::ptr::NonNull;
use std::{ptr, str};

use crate::result::*;

#[allow(missing_debug_implementations, missing_copy_implementations)]
pub struct RawConnection {
    internal_connection: NonNull<PGconn>,
}

impl RawConnection {
    pub fn establish(database_url: &str) -> ConnectionResult<Self> {
        use self::ConnStatusType::*;

        let connection_string = CString::new(database_url)?;
        let connection_ptr = unsafe { PQconnectdb(connection_string.as_ptr()) };
        let connection_status = unsafe { PQstatus(connection_ptr) };

        match connection_status {
            CONNECTION_OK => {
                let connection_ptr = unsafe { NonNull::new_unchecked(connection_ptr) };
                Ok(RawConnection {
                    internal_connection: connection_ptr,
                })
            }
            _ => {
                let message = last_error_message(connection_ptr);

                unsafe { PQfinish(connection_ptr) }

                Err(ConnectionError::BadConnection(message))
            }
        }
    }

    pub fn last_error_message(&self) -> String {
        last_error_message(self.internal_connection.as_ptr())
    }

    pub fn set_notice_processor(&self, notice_processor: NoticeProcessor) {
        unsafe {
            PQsetNoticeProcessor(
                self.internal_connection.as_ptr(),
                Some(notice_processor),
                ptr::null_mut(),
            );
        }
    }

    pub unsafe fn exec(&self, query: *const libc::c_char) -> QueryResult<RawResult> {
        RawResult::new(PQexec(self.internal_connection.as_ptr(), query), self)
    }

    pub unsafe fn exec_prepared(
        &self,
        stmt_name: *const libc::c_char,
        param_count: libc::c_int,
        param_values: *const *const libc::c_char,
        param_lengths: *const libc::c_int,
        param_formats: *const libc::c_int,
        result_format: libc::c_int,
    ) -> QueryResult<RawResult> {
        let ptr = PQexecPrepared(
            self.internal_connection.as_ptr(),
            stmt_name,
            param_count,
            param_values,
            param_lengths,
            param_formats,
            result_format,
        );
        RawResult::new(ptr, self)
    }

    pub unsafe fn prepare(
        &self,
        stmt_name: *const libc::c_char,
        query: *const libc::c_char,
        param_count: libc::c_int,
        param_types: *const Oid,
    ) -> QueryResult<RawResult> {
        let ptr = PQprepare(
            self.internal_connection.as_ptr(),
            stmt_name,
            query,
            param_count,
            param_types,
        );
        RawResult::new(ptr, self)
    }
}

pub type NoticeProcessor = extern "C" fn(arg: *mut libc::c_void, message: *const libc::c_char);

impl Drop for RawConnection {
    fn drop(&mut self) {
        unsafe { PQfinish(self.internal_connection.as_ptr()) };
    }
}

fn last_error_message(conn: *const PGconn) -> String {
    unsafe {
        let error_ptr = PQerrorMessage(conn);
        let bytes = CStr::from_ptr(error_ptr).to_bytes();
        String::from_utf8_lossy(bytes).to_string()
    }
}

/// Internal wrapper around a `*mut PGresult` which is known to be not-null, and
/// have no aliases.  This wrapper is to ensure that it's always properly
/// dropped.
///
/// If `Unique` is ever stabilized, we should use it here.
#[allow(missing_debug_implementations)]
pub struct RawResult(NonNull<PGresult>);

unsafe impl Send for RawResult {}
unsafe impl Sync for RawResult {}

impl RawResult {
    #[allow(clippy::new_ret_no_self)]
    fn new(ptr: *mut PGresult, conn: &RawConnection) -> QueryResult<Self> {
        NonNull::new(ptr).map(RawResult).ok_or_else(|| {
            Error::DatabaseError(
                DatabaseErrorKind::UnableToSendCommand,
                Box::new(conn.last_error_message()),
            )
        })
    }

    pub fn as_ptr(&self) -> *mut PGresult {
        self.0.as_ptr()
    }

    pub fn error_message(&self) -> &str {
        let ptr = unsafe { PQresultErrorMessage(self.0.as_ptr()) };
        let cstr = unsafe { CStr::from_ptr(ptr) };
        cstr.to_str().unwrap_or_default()
    }
}

impl Drop for RawResult {
    fn drop(&mut self) {
        unsafe { PQclear(self.0.as_ptr()) }
    }
}
<commit_msg>Check that the connection pointer is not null before trying to call PQfinish after a PostgreSQL connection attempt failure<commit_after>#![allow(clippy::too_many_arguments)]

extern crate pq_sys;

use self::pq_sys::*;
use std::ffi::{CStr, CString};
use std::os::raw as libc;
use std::ptr::NonNull;
use std::{ptr, str};

use crate::result::*;

#[allow(missing_debug_implementations, missing_copy_implementations)]
pub struct RawConnection {
    internal_connection: NonNull<PGconn>,
}

impl RawConnection {
    pub fn establish(database_url: &str) -> ConnectionResult<Self> {
        use self::ConnStatusType::*;

        let connection_string = CString::new(database_url)?;
        let connection_ptr = unsafe { PQconnectdb(connection_string.as_ptr()) };
        let connection_status = unsafe { PQstatus(connection_ptr) };

        match connection_status {
            CONNECTION_OK => {
                let connection_ptr = unsafe { NonNull::new_unchecked(connection_ptr) };
                Ok(RawConnection {
                    internal_connection: connection_ptr,
                })
            }
            _ => {
                let message = last_error_message(connection_ptr);

                if !connection_ptr.is_null() {
                     // Note that even if the server connection attempt fails (as indicated by PQstatus),
                     // the application should call PQfinish to free the memory used by the PGconn object.
                     // https://www.postgresql.org/docs/current/libpq-connect.html
                     unsafe { PQfinish(connection_ptr) }
                }

                Err(ConnectionError::BadConnection(message))
            }
        }
    }

    pub fn last_error_message(&self) -> String {
        last_error_message(self.internal_connection.as_ptr())
    }

    pub fn set_notice_processor(&self, notice_processor: NoticeProcessor) {
        unsafe {
            PQsetNoticeProcessor(
                self.internal_connection.as_ptr(),
                Some(notice_processor),
                ptr::null_mut(),
            );
        }
    }

    pub unsafe fn exec(&self, query: *const libc::c_char) -> QueryResult<RawResult> {
        RawResult::new(PQexec(self.internal_connection.as_ptr(), query), self)
    }

    pub unsafe fn exec_prepared(
        &self,
        stmt_name: *const libc::c_char,
        param_count: libc::c_int,
        param_values: *const *const libc::c_char,
        param_lengths: *const libc::c_int,
        param_formats: *const libc::c_int,
        result_format: libc::c_int,
    ) -> QueryResult<RawResult> {
        let ptr = PQexecPrepared(
            self.internal_connection.as_ptr(),
            stmt_name,
            param_count,
            param_values,
            param_lengths,
            param_formats,
            result_format,
        );
        RawResult::new(ptr, self)
    }

    pub unsafe fn prepare(
        &self,
        stmt_name: *const libc::c_char,
        query: *const libc::c_char,
        param_count: libc::c_int,
        param_types: *const Oid,
    ) -> QueryResult<RawResult> {
        let ptr = PQprepare(
            self.internal_connection.as_ptr(),
            stmt_name,
            query,
            param_count,
            param_types,
        );
        RawResult::new(ptr, self)
    }
}

pub type NoticeProcessor = extern "C" fn(arg: *mut libc::c_void, message: *const libc::c_char);

impl Drop for RawConnection {
    fn drop(&mut self) {
        unsafe { PQfinish(self.internal_connection.as_ptr()) };
    }
}

fn last_error_message(conn: *const PGconn) -> String {
    unsafe {
        let error_ptr = PQerrorMessage(conn);
        let bytes = CStr::from_ptr(error_ptr).to_bytes();
        String::from_utf8_lossy(bytes).to_string()
    }
}

/// Internal wrapper around a `*mut PGresult` which is known to be not-null, and
/// have no aliases.  This wrapper is to ensure that it's always properly
/// dropped.
///
/// If `Unique` is ever stabilized, we should use it here.
#[allow(missing_debug_implementations)]
pub struct RawResult(NonNull<PGresult>);

unsafe impl Send for RawResult {}
unsafe impl Sync for RawResult {}

impl RawResult {
    #[allow(clippy::new_ret_no_self)]
    fn new(ptr: *mut PGresult, conn: &RawConnection) -> QueryResult<Self> {
        NonNull::new(ptr).map(RawResult).ok_or_else(|| {
            Error::DatabaseError(
                DatabaseErrorKind::UnableToSendCommand,
                Box::new(conn.last_error_message()),
            )
        })
    }

    pub fn as_ptr(&self) -> *mut PGresult {
        self.0.as_ptr()
    }

    pub fn error_message(&self) -> &str {
        let ptr = unsafe { PQresultErrorMessage(self.0.as_ptr()) };
        let cstr = unsafe { CStr::from_ptr(ptr) };
        cstr.to_str().unwrap_or_default()
    }
}

impl Drop for RawResult {
    fn drop(&mut self) {
        unsafe { PQclear(self.0.as_ptr()) }
    }
}
<|endoftext|>


<commit_before>
// NB (Morten, 09.10.17)
// Currently, we assume SampleData to be i16!
// See SND_PCM_FORMAT_S16_LE

extern crate alsa_sys as alsa;

use std::mem;
use std::ffi::CStr;

use super::*;
use time::Time;

const MAX_WRITE_FRAMES: u64 = 1024;

pub(super) struct AudioBackend {
    pcm_handle: *mut alsa::snd_pcm_t,
    write_buffer: Vec<i16>,
    total_frames: u64,
}

impl AudioBackend {
    pub fn initialize() -> Result<AudioBackend, ()> {
        let mut pcm_handle = ptr::null_mut();
        let mut write_buffer = Vec::new();
        let total_frames;

        unsafe {
            let device_name = b"default\0";

            let result = alsa::snd_pcm_open(
                &mut pcm_handle,
                device_name.as_ptr() as *const i8,
                alsa::SND_PCM_STREAM_PLAYBACK,
                0
            );
            if result < 0 {
                println!("snd_pcm_open failed: {}", result);
                return Err(());
            }

            // Configure "hardware" stuff
            let mut hardware = ptr::null_mut();
            let result = alsa::snd_pcm_hw_params_malloc(&mut hardware);
            if result < 0 {
                println!("snd_pcm_hw_params_malloc failed: {}", result);
                return Err(());
            }
            assert!(!hardware.is_null());

            let result = alsa::snd_pcm_hw_params_any(pcm_handle, hardware);
            if result < 0 {
                println!("snd_pcm_hw_params_any failed: {}", result);
                return Err(());
            }

            let access = alsa::SND_PCM_ACCESS_RW_INTERLEAVED;
            let format = if cfg!(target_endian = "big") {
                alsa::SND_PCM_FORMAT_S16_BE
            } else {
                alsa::SND_PCM_FORMAT_S16_LE
            };
            let channels = OUTPUT_CHANNELS;
            let mut sample_rate = OUTPUT_SAMPLE_RATE;

            alsa::snd_pcm_hw_params_set_access(pcm_handle, hardware, access);
            alsa::snd_pcm_hw_params_set_format(pcm_handle, hardware, format);
            alsa::snd_pcm_hw_params_set_channels(pcm_handle, hardware, channels);
            alsa::snd_pcm_hw_params_set_rate_near(pcm_handle, hardware, &mut sample_rate, ptr::null_mut());

            let result = alsa::snd_pcm_hw_params(pcm_handle, hardware);
            if result < 0 {
                println!("snd_pcm_hw_params failed: {}", result);
                return Err(());
            }

            alsa::snd_pcm_hw_params_free(hardware);

            // Configure "software" stuff
            let mut software = ptr::null_mut();
            let result = alsa::snd_pcm_sw_params_malloc(&mut software);
            if result < 0 {
                println!("snd_pcm_sw_params_malloc failed: {}", result);
                return Err(());
            }
            assert!(!software.is_null());

            let result = alsa::snd_pcm_sw_params_current(pcm_handle, software);
            if result < 0 {
                println!("snd_pcm_sw_params_current failed: {}", result);
                return Err(());
            }

            alsa::snd_pcm_sw_params_set_avail_min(pcm_handle, software, MAX_WRITE_FRAMES);
            alsa::snd_pcm_sw_params_set_start_threshold(pcm_handle, software, 0);

            let result = alsa::snd_pcm_sw_params(pcm_handle, software);
            if result < 0 {
                println!("snd_pcm_sw_params failed: {}", result);
                return Err(());
            }

            alsa::snd_pcm_sw_params_free(software);

            total_frames = alsa::snd_pcm_avail(pcm_handle) as u64;

            // Prepare for playing
            let result = alsa::snd_pcm_prepare(pcm_handle);
            if result < 0 {
                println!("snd_pcm_prepare failed: {}", result);
                return Err(());
            }

            // Write some bytes at the start to prevent buffer underruns
            let samples = MAX_WRITE_FRAMES as usize * OUTPUT_CHANNELS as usize;
            write_buffer.reserve(samples);
            ptr::write_bytes(write_buffer.as_mut_ptr(), 0, samples);

            let result = alsa::snd_pcm_writei(
                pcm_handle,
                write_buffer.as_ptr() as *const _,
                MAX_WRITE_FRAMES,
            );
            if result < 0 {
                println!("snd_pcm_writei failed: {}", result);
                return Err(());
            }
        }

        Ok(AudioBackend {
            pcm_handle,
            write_buffer,
            total_frames,
        })
    }

    pub fn write<F>(
        &mut self,
        frame_counter: &mut u64,
        mut mix_callback: F,
    ) -> Result<bool, ()>
      where F: FnMut(u64, &mut [SampleData]),
    {
        // ALSA will request enough frames to fill up the entire ring buffer,
        // we only want to write a few frames ahead to keep latency low.

        let available_frames;

        unsafe {
            let result = alsa::snd_pcm_avail_update(self.pcm_handle);
            if result < 0 {
                if result == -32 {
                    // underrun, we did not provide data quickly enough. What do we do now?
                    // TODO in debug mode this happens when mix_callback has many sounds. We
                    // probably want to "gracefully" handle that case, by stuttering
                }

                println!("snd_pcm_avail_delay failed: {}", result);
                return Err(());
            } else {
                available_frames = result as u64;
            }
        }

        if available_frames <= 0 {
            // We somehow managed to fill up the entire ring buffer, this is sort of bad
            return Ok(false);
        }

        let unplayed_frames = self.total_frames - available_frames;
        if unplayed_frames > 2*MAX_WRITE_FRAMES {
            return Ok(false);
        }

        let write_frames = if unplayed_frames < MAX_WRITE_FRAMES {
            2*MAX_WRITE_FRAMES - unplayed_frames
        } else {
            MAX_WRITE_FRAMES
        };
        let samples = write_frames as usize * OUTPUT_CHANNELS as usize;

        self.write_buffer.clear();
        self.write_buffer.reserve(samples);
        unsafe {
            self.write_buffer.set_len(samples);
            ptr::write_bytes(self.write_buffer.as_mut_ptr(), 0, samples);
        }

        mix_callback(*frame_counter, &mut self.write_buffer);
        *frame_counter += write_frames;

        unsafe {
            let result = alsa::snd_pcm_writei(
                self.pcm_handle,
                self.write_buffer.as_ptr() as *const _,
                write_frames,
            );
            if result < 0 {
                println!("snd_pcm_writei failed: {}", result);
                return Err(());
            }
        }

        return Ok(true); // We wrote some data
    }

    pub fn write_interval(&self) -> Time {
        Time((MAX_WRITE_FRAMES as u64 * Time::NANOSECONDS_PER_SECOND) / OUTPUT_SAMPLE_RATE as u64)
    }
}

impl Drop for AudioBackend {
    fn drop(&mut self) {
        unsafe {
            alsa::snd_pcm_close(self.pcm_handle);
        }
    }
}
<commit_msg>Partial underrun detection and recovery for ALSA<commit_after>
// NB (Morten, 09.10.17)
// Currently, we assume SampleData to be i16!
// See SND_PCM_FORMAT_S16_LE

extern crate alsa_sys as alsa;

use std::mem;
use std::ffi::CStr;

use super::*;
use time::Time;

const MAX_WRITE_FRAMES: u64 = 1024;

pub(super) struct AudioBackend {
    pcm_handle: *mut alsa::snd_pcm_t,
    write_buffer: Vec<i16>,
    total_frames: u64,
}

impl AudioBackend {
    pub fn initialize() -> Result<AudioBackend, ()> {
        let mut pcm_handle = ptr::null_mut();
        let mut write_buffer = Vec::new();
        let total_frames;

        unsafe {
            let device_name = b"default\0";

            let result = alsa::snd_pcm_open(
                &mut pcm_handle,
                device_name.as_ptr() as *const i8,
                alsa::SND_PCM_STREAM_PLAYBACK,
                0
            );
            if result < 0 {
                println!("snd_pcm_open failed: {}", result);
                return Err(());
            }

            // Configure "hardware" stuff
            let mut hardware = ptr::null_mut();
            let result = alsa::snd_pcm_hw_params_malloc(&mut hardware);
            if result < 0 {
                println!("snd_pcm_hw_params_malloc failed: {}", result);
                return Err(());
            }
            assert!(!hardware.is_null());

            let result = alsa::snd_pcm_hw_params_any(pcm_handle, hardware);
            if result < 0 {
                println!("snd_pcm_hw_params_any failed: {}", result);
                return Err(());
            }

            let access = alsa::SND_PCM_ACCESS_RW_INTERLEAVED;
            let format = if cfg!(target_endian = "big") {
                alsa::SND_PCM_FORMAT_S16_BE
            } else {
                alsa::SND_PCM_FORMAT_S16_LE
            };
            let channels = OUTPUT_CHANNELS;
            let mut sample_rate = OUTPUT_SAMPLE_RATE;

            alsa::snd_pcm_hw_params_set_access(pcm_handle, hardware, access);
            alsa::snd_pcm_hw_params_set_format(pcm_handle, hardware, format);
            alsa::snd_pcm_hw_params_set_channels(pcm_handle, hardware, channels);
            alsa::snd_pcm_hw_params_set_rate_near(pcm_handle, hardware, &mut sample_rate, ptr::null_mut());

            let result = alsa::snd_pcm_hw_params(pcm_handle, hardware);
            if result < 0 {
                println!("snd_pcm_hw_params failed: {}", result);
                return Err(());
            }

            alsa::snd_pcm_hw_params_free(hardware);

            // Configure "software" stuff
            let mut software = ptr::null_mut();
            let result = alsa::snd_pcm_sw_params_malloc(&mut software);
            if result < 0 {
                println!("snd_pcm_sw_params_malloc failed: {}", result);
                return Err(());
            }
            assert!(!software.is_null());

            let result = alsa::snd_pcm_sw_params_current(pcm_handle, software);
            if result < 0 {
                println!("snd_pcm_sw_params_current failed: {}", result);
                return Err(());
            }

            alsa::snd_pcm_sw_params_set_avail_min(pcm_handle, software, MAX_WRITE_FRAMES);
            alsa::snd_pcm_sw_params_set_start_threshold(pcm_handle, software, 0);

            let result = alsa::snd_pcm_sw_params(pcm_handle, software);
            if result < 0 {
                println!("snd_pcm_sw_params failed: {}", result);
                return Err(());
            }

            alsa::snd_pcm_sw_params_free(software);

            total_frames = alsa::snd_pcm_avail(pcm_handle) as u64;

            // Prepare for playing
            let result = alsa::snd_pcm_prepare(pcm_handle);
            if result < 0 {
                println!("snd_pcm_prepare failed: {}", result);
                return Err(());
            }

            // Write some bytes at the start to prevent buffer underruns
            let samples = MAX_WRITE_FRAMES as usize * OUTPUT_CHANNELS as usize;
            write_buffer.reserve(samples);
            ptr::write_bytes(write_buffer.as_mut_ptr(), 0, samples);

            let result = alsa::snd_pcm_writei(
                pcm_handle,
                write_buffer.as_ptr() as *const _,
                MAX_WRITE_FRAMES,
            );
            if result < 0 {
                println!("snd_pcm_writei failed: {}", result);
                return Err(());
            }
        }

        Ok(AudioBackend {
            pcm_handle,
            write_buffer,
            total_frames,
        })
    }

    pub fn write<F>(
        &mut self,
        frame_counter: &mut u64,
        mut mix_callback: F,
    ) -> Result<bool, ()>
      where F: FnMut(u64, &mut [SampleData]),
    {
        // ALSA will request enough frames to fill up the entire ring buffer,
        // we only want to write a few frames ahead to keep latency low.

        let available_frames;

        unsafe {
            let result = alsa::snd_pcm_avail_update(self.pcm_handle);
            if result == -32 {
                // We did not provide data fast enough, recover
                let recover_result = alsa::snd_pcm_recover(self.pcm_handle, -32, 1);
                if recover_result < 0 {
                    println!("Underrun detected, could not recover");
                    return Err(()); // We are probably fucked
                } else {

                    // Try again
                    let retry_result = alsa::snd_pcm_avail_update(self.pcm_handle);
                    if retry_result < 0 {
                        println!("Underrun detected, recovered but it did not help");
                        return Err(());
                    } else {
                        println!("Underrun detected and fixed");
                        available_frames = retry_result as u64;
                    }
                }

            } else if result < 0 {
                println!("snd_pcm_avail_delay failed: {}", result);
                return Err(());
            } else {
                available_frames = result as u64;
            }
        }

        if available_frames <= 0 {
            // We somehow managed to fill up the entire ring buffer, this is sort of bad
            return Ok(false);
        }

        let unplayed_frames = self.total_frames - available_frames;
        if unplayed_frames > 2*MAX_WRITE_FRAMES {
            return Ok(false);
        }

        let write_frames = if unplayed_frames < MAX_WRITE_FRAMES {
            2*MAX_WRITE_FRAMES - unplayed_frames
        } else {
            MAX_WRITE_FRAMES
        };
        let samples = write_frames as usize * OUTPUT_CHANNELS as usize;

        self.write_buffer.clear();
        self.write_buffer.reserve(samples);
        unsafe {
            self.write_buffer.set_len(samples);
            ptr::write_bytes(self.write_buffer.as_mut_ptr(), 0, samples);
        }

        mix_callback(*frame_counter, &mut self.write_buffer);
        *frame_counter += write_frames;

        unsafe {
            // TODO we might also get a underrun here, we probably can recover from that as well!
            let result = alsa::snd_pcm_writei(
                self.pcm_handle,
                self.write_buffer.as_ptr() as *const _,
                write_frames,
            );
            if result < 0 {
                println!("snd_pcm_writei failed: {}", result);
                return Err(());
            }
        }

        return Ok(true); // We wrote some data
    }

    pub fn write_interval(&self) -> Time {
        Time((MAX_WRITE_FRAMES as u64 * Time::NANOSECONDS_PER_SECOND) / OUTPUT_SAMPLE_RATE as u64)
    }
}

impl Drop for AudioBackend {
    fn drop(&mut self) {
        unsafe {
            alsa::snd_pcm_close(self.pcm_handle);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use std::collections::BTreeMap;

use anyhow::anyhow;
use itertools::zip_eq;

use consensus_types::{
    common::Author, executed_block::ExecutedBlock, experimental::commit_vote::CommitVote,
};
use diem_crypto::ed25519::Ed25519Signature;
use diem_types::{
    account_address::AccountAddress,
    block_info::BlockInfo,
    ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
    validator_verifier::ValidatorVerifier,
};

use crate::state_replication::StateComputerCommitCallBackType;
use diem_crypto::HashValue;

// we differentiate buffer items at different stages
// for better code readability
pub struct OrderedBufferItem {
    pub pending_votes: BTreeMap<AccountAddress, Ed25519Signature>,
    pub callback: StateComputerCommitCallBackType,
    pub ordered_blocks: Vec<ExecutedBlock>,
    pub ordered_proof: LedgerInfoWithSignatures,
}

pub struct ExecutedBufferItem {
    pub executed_blocks: Vec<ExecutedBlock>,
    pub pending_votes: BTreeMap<AccountAddress, Ed25519Signature>,
    pub callback: StateComputerCommitCallBackType,
    pub commit_info: BlockInfo,
    pub ordered_proof: LedgerInfoWithSignatures,
}

impl ExecutedBufferItem {
    pub fn generate_commit_ledger_info(&self) -> LedgerInfo {
        LedgerInfo::new(
            self.commit_info.clone(),
            self.ordered_proof.ledger_info().consensus_data_hash(),
        )
    }
}

pub struct SignedBufferItem {
    pub executed_blocks: Vec<ExecutedBlock>,
    pub commit_proof: LedgerInfoWithSignatures,
    pub callback: StateComputerCommitCallBackType,
    pub commit_vote: CommitVote,
}

pub struct AggregatedBufferItem {
    pub executed_blocks: Vec<ExecutedBlock>,
    pub aggregated_proof: LedgerInfoWithSignatures,
    pub callback: StateComputerCommitCallBackType,
}

pub enum BufferItem {
    Ordered(Box<OrderedBufferItem>),
    Executed(Box<ExecutedBufferItem>),
    Signed(Box<SignedBufferItem>),
    Aggregated(Box<AggregatedBufferItem>),
}

impl BufferItem {
    pub fn new_ordered(
        ordered_blocks: Vec<ExecutedBlock>,
        ordered_proof: LedgerInfoWithSignatures,
        callback: StateComputerCommitCallBackType,
    ) -> Self {
        Self::Ordered(Box::new(OrderedBufferItem {
            pending_votes: BTreeMap::<AccountAddress, Ed25519Signature>::new(),
            callback,
            ordered_blocks,
            ordered_proof,
        }))
    }

    // pipeline functions
    pub fn advance_to_executed(self, executed_blocks: Vec<ExecutedBlock>) -> Self {
        match self {
            Self::Ordered(ordered_item) => {
                let OrderedBufferItem {
                    ordered_blocks,
                    pending_votes,
                    callback,
                    ordered_proof,
                } = *ordered_item;
                for (b1, b2) in zip_eq(ordered_blocks.iter(), executed_blocks.iter()) {
                    assert_eq!(b1.id(), b2.id());
                }
                let commit_info = executed_blocks.last().unwrap().block_info();
                Self::Executed(Box::new(ExecutedBufferItem {
                    executed_blocks,
                    pending_votes,
                    callback,
                    commit_info,
                    ordered_proof,
                }))
            }
            _ => {
                panic!("Only ordered blocks can advance to executed blocks.")
            }
        }
    }

    pub fn advance_to_signed(
        self,
        author: Author,
        signature: Ed25519Signature,
        verifier: &ValidatorVerifier,
    ) -> (Self, CommitVote) {
        match self {
            Self::Executed(executed_item) => {
                let commit_ledger_info = executed_item.generate_commit_ledger_info();
                let ExecutedBufferItem {
                    executed_blocks,
                    callback,
                    pending_votes,
                    ..
                } = *executed_item;

                let mut valid_sigs = BTreeMap::<AccountAddress, Ed25519Signature>::new();
                valid_sigs.insert(author, signature.clone());

                for (author, sig) in pending_votes.iter() {
                    if verifier.verify(*author, &commit_ledger_info, sig).is_ok() {
                        valid_sigs.insert(*author, sig.clone());
                    }
                }

                let commit_ledger_info_with_sigs =
                    LedgerInfoWithSignatures::new(commit_ledger_info.clone(), valid_sigs);

                let commit_vote =
                    CommitVote::new_with_signature(author, commit_ledger_info, signature);

                (
                    Self::Signed(Box::new(SignedBufferItem {
                        executed_blocks,
                        callback,
                        commit_proof: commit_ledger_info_with_sigs,
                        commit_vote: commit_vote.clone(),
                    })),
                    commit_vote,
                )
            }
            _ => {
                panic!("Only executed buffer items can advance to signed blocks.")
            }
        }
    }

    /// this function assumes block id matches and the validity of ledger_info and that it has the voting power
    /// it returns an updated item
    pub fn try_advance_to_aggregated_with_ledger_info(
        self,
        commit_ledger_info: LedgerInfoWithSignatures,
    ) -> Self {
        match self {
            Self::Signed(signed_item) => {
                let SignedBufferItem {
                    executed_blocks,
                    callback,
                    commit_proof,
                    ..
                } = *signed_item;
                assert_eq!(commit_proof.commit_info(), commit_ledger_info.commit_info(),);
                Self::Aggregated(Box::new(AggregatedBufferItem {
                    executed_blocks,
                    callback,
                    aggregated_proof: commit_ledger_info,
                }))
            }
            Self::Executed(executed_item) => {
                let ExecutedBufferItem {
                    executed_blocks,
                    callback,
                    commit_info,
                    ..
                } = *executed_item;
                assert_eq!(commit_info, *commit_ledger_info.commit_info());
                Self::Aggregated(Box::new(AggregatedBufferItem {
                    executed_blocks,
                    callback,
                    aggregated_proof: commit_ledger_info,
                }))
            }
            Self::Ordered(ordered_item) => {
                let ordered = *ordered_item;
                assert!(ordered
                    .ordered_proof
                    .commit_info()
                    .match_ordered_only(commit_ledger_info.commit_info()));
                Self::Ordered(Box::new(OrderedBufferItem {
                    pending_votes: commit_ledger_info.signatures().clone(),
                    ..ordered
                }))
            }
            Self::Aggregated(_) => {
                unreachable!("Found aggregated buffer item but any aggregated buffer item should get dequeued right away.");
            }
        }
    }

    pub fn try_advance_to_aggregated(self, validator: &ValidatorVerifier) -> Self {
        match self {
            Self::Signed(signed_item) => {
                if signed_item
                    .commit_proof
                    .check_voting_power(validator)
                    .is_ok()
                {
                    Self::Aggregated(Box::new(AggregatedBufferItem {
                        executed_blocks: signed_item.executed_blocks,
                        aggregated_proof: signed_item.commit_proof,
                        callback: signed_item.callback,
                    }))
                } else {
                    Self::Signed(signed_item)
                }
            }
            Self::Executed(executed_item) => {
                // TODO this needs to verify the votes match the commit info
                if validator
                    .check_voting_power(executed_item.pending_votes.keys())
                    .is_ok()
                {
                    let aggregated_proof = LedgerInfoWithSignatures::new(
                        executed_item.generate_commit_ledger_info(),
                        executed_item.pending_votes,
                    );
                    Self::Aggregated(Box::new(AggregatedBufferItem {
                        executed_blocks: executed_item.executed_blocks,
                        aggregated_proof,
                        callback: executed_item.callback,
                    }))
                } else {
                    Self::Executed(executed_item)
                }
            }
            _ => self,
        }
    }

    // generic functions
    pub fn get_blocks(&self) -> &Vec<ExecutedBlock> {
        match self {
            Self::Ordered(ordered) => &ordered.ordered_blocks,
            Self::Executed(executed) => &executed.executed_blocks,
            Self::Signed(signed) => &signed.executed_blocks,
            Self::Aggregated(aggregated) => &aggregated.executed_blocks,
        }
    }

    pub fn block_id(&self) -> HashValue {
        self.get_blocks().last().unwrap().id()
    }

    pub fn get_commit_info(&self) -> &BlockInfo {
        match self {
            Self::Ordered(_) => {
                panic!("Ordered buffer item does not contain commit info");
            }
            Self::Executed(executed) => &executed.commit_info,
            Self::Signed(signed) => signed.commit_proof.ledger_info().commit_info(),
            Self::Aggregated(aggregated) => aggregated.aggregated_proof.ledger_info().commit_info(),
        }
    }

    pub fn add_signature_if_matched(
        &mut self,
        target_commit_info: &BlockInfo,
        author: Author,
        sig: Ed25519Signature,
    ) -> anyhow::Result<()> {
        match self {
            Self::Ordered(ordered) => {
                if ordered
                    .ordered_proof
                    .commit_info()
                    .match_ordered_only(target_commit_info)
                {
                    // we optimistically assume the vote will be valid in the future.
                    // when advancing to signed item, we will check if the sigs are valid.
                    // each author at most stores a single sig for each item,
                    // so an adversary will not be able to flood our memory.
                    ordered.pending_votes.insert(author, sig);
                    return Ok(());
                }
            }
            Self::Executed(executed) => {
                if &executed.commit_info == target_commit_info {
                    executed.pending_votes.insert(author, sig);
                    return Ok(());
                }
            }
            Self::Signed(signed) => {
                if signed.commit_proof.commit_info() == target_commit_info {
                    signed.commit_proof.add_signature(author, sig);
                    return Ok(());
                }
            }
            Self::Aggregated(aggregated) => {
                // we do not need to do anything for aggregated
                // but return true is helpful to stop the outer loop early
                if aggregated.aggregated_proof.commit_info() == target_commit_info {
                    return Ok(());
                }
            }
        }
        Err(anyhow!("Inconsistent commit info."))
    }

    pub fn has_been_executed(&self) -> bool {
        !self.is_ordered()
    }

    pub fn is_ordered(&self) -> bool {
        matches!(self, Self::Ordered(_))
    }

    pub fn is_executed(&self) -> bool {
        matches!(self, Self::Executed(_))
    }

    pub fn is_signed(&self) -> bool {
        matches!(self, Self::Signed(_))
    }

    pub fn is_aggregated(&self) -> bool {
        matches!(self, Self::Aggregated(_))
    }
}
<commit_msg>[decoupled-execution] verify votes in try_advance_to_aggregated<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use std::collections::BTreeMap;

use anyhow::anyhow;
use itertools::zip_eq;

use consensus_types::{
    common::Author, executed_block::ExecutedBlock, experimental::commit_vote::CommitVote,
};
use diem_crypto::ed25519::Ed25519Signature;
use diem_types::{
    account_address::AccountAddress,
    block_info::BlockInfo,
    ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
    validator_verifier::ValidatorVerifier,
};

use crate::state_replication::StateComputerCommitCallBackType;
use diem_crypto::HashValue;

// we differentiate buffer items at different stages
// for better code readability
pub struct OrderedBufferItem {
    pub pending_votes: BTreeMap<AccountAddress, Ed25519Signature>,
    pub callback: StateComputerCommitCallBackType,
    pub ordered_blocks: Vec<ExecutedBlock>,
    pub ordered_proof: LedgerInfoWithSignatures,
}

pub struct ExecutedBufferItem {
    pub executed_blocks: Vec<ExecutedBlock>,
    pub pending_votes: BTreeMap<AccountAddress, Ed25519Signature>,
    pub callback: StateComputerCommitCallBackType,
    pub commit_info: BlockInfo,
    pub ordered_proof: LedgerInfoWithSignatures,
}

impl ExecutedBufferItem {
    pub fn generate_commit_ledger_info(&self) -> LedgerInfo {
        LedgerInfo::new(
            self.commit_info.clone(),
            self.ordered_proof.ledger_info().consensus_data_hash(),
        )
    }
}

pub struct SignedBufferItem {
    pub executed_blocks: Vec<ExecutedBlock>,
    pub commit_proof: LedgerInfoWithSignatures,
    pub callback: StateComputerCommitCallBackType,
    pub commit_vote: CommitVote,
}

pub struct AggregatedBufferItem {
    pub executed_blocks: Vec<ExecutedBlock>,
    pub aggregated_proof: LedgerInfoWithSignatures,
    pub callback: StateComputerCommitCallBackType,
}

pub enum BufferItem {
    Ordered(Box<OrderedBufferItem>),
    Executed(Box<ExecutedBufferItem>),
    Signed(Box<SignedBufferItem>),
    Aggregated(Box<AggregatedBufferItem>),
}

impl BufferItem {
    pub fn new_ordered(
        ordered_blocks: Vec<ExecutedBlock>,
        ordered_proof: LedgerInfoWithSignatures,
        callback: StateComputerCommitCallBackType,
    ) -> Self {
        Self::Ordered(Box::new(OrderedBufferItem {
            pending_votes: BTreeMap::<AccountAddress, Ed25519Signature>::new(),
            callback,
            ordered_blocks,
            ordered_proof,
        }))
    }

    // pipeline functions
    pub fn advance_to_executed(self, executed_blocks: Vec<ExecutedBlock>) -> Self {
        match self {
            Self::Ordered(ordered_item) => {
                let OrderedBufferItem {
                    ordered_blocks,
                    pending_votes,
                    callback,
                    ordered_proof,
                } = *ordered_item;
                for (b1, b2) in zip_eq(ordered_blocks.iter(), executed_blocks.iter()) {
                    assert_eq!(b1.id(), b2.id());
                }
                let commit_info = executed_blocks.last().unwrap().block_info();
                Self::Executed(Box::new(ExecutedBufferItem {
                    executed_blocks,
                    pending_votes,
                    callback,
                    commit_info,
                    ordered_proof,
                }))
            }
            _ => {
                panic!("Only ordered blocks can advance to executed blocks.")
            }
        }
    }

    fn aggregate_ledger_info(
        commit_ledger_info: &LedgerInfo,
        signatures: BTreeMap<AccountAddress, Ed25519Signature>,
        validator: &ValidatorVerifier,
    ) -> LedgerInfoWithSignatures {
        let mut valid_sigs = BTreeMap::new();

        for (author, sig) in signatures.into_iter() {
            if validator.verify(author, commit_ledger_info, &sig).is_ok() {
                valid_sigs.insert(author, sig);
            }
        }

        LedgerInfoWithSignatures::new(commit_ledger_info.clone(), valid_sigs)
    }

    pub fn advance_to_signed(
        self,
        author: Author,
        signature: Ed25519Signature,
        validator: &ValidatorVerifier,
    ) -> (Self, CommitVote) {
        match self {
            Self::Executed(executed_item) => {
                let commit_ledger_info = executed_item.generate_commit_ledger_info();
                let ExecutedBufferItem {
                    executed_blocks,
                    callback,
                    mut pending_votes,
                    ..
                } = *executed_item;

                pending_votes.insert(author, signature.clone());
                let commit_proof =
                    Self::aggregate_ledger_info(&commit_ledger_info, pending_votes, validator);

                let commit_vote =
                    CommitVote::new_with_signature(author, commit_ledger_info, signature);

                (
                    Self::Signed(Box::new(SignedBufferItem {
                        executed_blocks,
                        callback,
                        commit_proof,
                        commit_vote: commit_vote.clone(),
                    })),
                    commit_vote,
                )
            }
            _ => {
                panic!("Only executed buffer items can advance to signed blocks.")
            }
        }
    }

    /// this function assumes block id matches and the validity of ledger_info and that it has the voting power
    /// it returns an updated item
    pub fn try_advance_to_aggregated_with_ledger_info(
        self,
        commit_ledger_info: LedgerInfoWithSignatures,
    ) -> Self {
        match self {
            Self::Signed(signed_item) => {
                let SignedBufferItem {
                    executed_blocks,
                    callback,
                    commit_proof,
                    ..
                } = *signed_item;
                assert_eq!(commit_proof.commit_info(), commit_ledger_info.commit_info(),);
                Self::Aggregated(Box::new(AggregatedBufferItem {
                    executed_blocks,
                    callback,
                    aggregated_proof: commit_ledger_info,
                }))
            }
            Self::Executed(executed_item) => {
                let ExecutedBufferItem {
                    executed_blocks,
                    callback,
                    commit_info,
                    ..
                } = *executed_item;
                assert_eq!(commit_info, *commit_ledger_info.commit_info());
                Self::Aggregated(Box::new(AggregatedBufferItem {
                    executed_blocks,
                    callback,
                    aggregated_proof: commit_ledger_info,
                }))
            }
            Self::Ordered(ordered_item) => {
                let ordered = *ordered_item;
                assert!(ordered
                    .ordered_proof
                    .commit_info()
                    .match_ordered_only(commit_ledger_info.commit_info()));
                Self::Ordered(Box::new(OrderedBufferItem {
                    pending_votes: commit_ledger_info.signatures().clone(),
                    ..ordered
                }))
            }
            Self::Aggregated(_) => {
                unreachable!("Found aggregated buffer item but any aggregated buffer item should get dequeued right away.");
            }
        }
    }

    pub fn try_advance_to_aggregated(self, validator: &ValidatorVerifier) -> Self {
        match self {
            Self::Signed(signed_item) => {
                if signed_item
                    .commit_proof
                    .check_voting_power(validator)
                    .is_ok()
                {
                    Self::Aggregated(Box::new(AggregatedBufferItem {
                        executed_blocks: signed_item.executed_blocks,
                        aggregated_proof: signed_item.commit_proof,
                        callback: signed_item.callback,
                    }))
                } else {
                    Self::Signed(signed_item)
                }
            }
            Self::Executed(executed_item) => {
                let maybe_aggregated_proof = Self::aggregate_ledger_info(
                    &executed_item.generate_commit_ledger_info(),
                    executed_item.pending_votes.clone(),
                    validator,
                );
                if validator
                    .check_voting_power(maybe_aggregated_proof.signatures().keys())
                    .is_ok()
                {
                    Self::Aggregated(Box::new(AggregatedBufferItem {
                        executed_blocks: executed_item.executed_blocks,
                        aggregated_proof: maybe_aggregated_proof,
                        callback: executed_item.callback,
                    }))
                } else {
                    Self::Executed(executed_item)
                }
            }
            _ => self,
        }
    }

    // generic functions
    pub fn get_blocks(&self) -> &Vec<ExecutedBlock> {
        match self {
            Self::Ordered(ordered) => &ordered.ordered_blocks,
            Self::Executed(executed) => &executed.executed_blocks,
            Self::Signed(signed) => &signed.executed_blocks,
            Self::Aggregated(aggregated) => &aggregated.executed_blocks,
        }
    }

    pub fn block_id(&self) -> HashValue {
        self.get_blocks().last().unwrap().id()
    }

    pub fn get_commit_info(&self) -> &BlockInfo {
        match self {
            Self::Ordered(_) => {
                panic!("Ordered buffer item does not contain commit info");
            }
            Self::Executed(executed) => &executed.commit_info,
            Self::Signed(signed) => signed.commit_proof.ledger_info().commit_info(),
            Self::Aggregated(aggregated) => aggregated.aggregated_proof.ledger_info().commit_info(),
        }
    }

    pub fn add_signature_if_matched(
        &mut self,
        target_commit_info: &BlockInfo,
        author: Author,
        sig: Ed25519Signature,
    ) -> anyhow::Result<()> {
        match self {
            Self::Ordered(ordered) => {
                if ordered
                    .ordered_proof
                    .commit_info()
                    .match_ordered_only(target_commit_info)
                {
                    // we optimistically assume the vote will be valid in the future.
                    // when advancing to signed item, we will check if the sigs are valid.
                    // each author at most stores a single sig for each item,
                    // so an adversary will not be able to flood our memory.
                    ordered.pending_votes.insert(author, sig);
                    return Ok(());
                }
            }
            Self::Executed(executed) => {
                if &executed.commit_info == target_commit_info {
                    executed.pending_votes.insert(author, sig);
                    return Ok(());
                }
            }
            Self::Signed(signed) => {
                if signed.commit_proof.commit_info() == target_commit_info {
                    signed.commit_proof.add_signature(author, sig);
                    return Ok(());
                }
            }
            Self::Aggregated(aggregated) => {
                // we do not need to do anything for aggregated
                // but return true is helpful to stop the outer loop early
                if aggregated.aggregated_proof.commit_info() == target_commit_info {
                    return Ok(());
                }
            }
        }
        Err(anyhow!("Inconsistent commit info."))
    }

    pub fn has_been_executed(&self) -> bool {
        !self.is_ordered()
    }

    pub fn is_ordered(&self) -> bool {
        matches!(self, Self::Ordered(_))
    }

    pub fn is_executed(&self) -> bool {
        matches!(self, Self::Executed(_))
    }

    pub fn is_signed(&self) -> bool {
        matches!(self, Self::Signed(_))
    }

    pub fn is_aggregated(&self) -> bool {
        matches!(self, Self::Aggregated(_))
    }
}
<|endoftext|>


<commit_before>macro_rules! impl_raw_accessors(
    ($(($t:ty, $raw:ty)),+) => (
        $(
        impl $t {
            #[inline]
            pub unsafe fn raw(&self) -> $raw { self.raw }
        }
        )+
    )
);

macro_rules! impl_owned_accessors(
    ($(($t:ty, $owned:ident)),+) => (
        $(
        impl $t {
            #[inline]
            pub unsafe fn $owned(&self) -> bool { self.$owned }
        }
        )+
    )
);

macro_rules! impl_raw_constructor(
    ($(($t:ty, $te:ident ($($r:ident:$rt:ty),+))),+) => (
        $(
        impl $t {
            #[inline]
            pub unsafe fn from_ll($($r:$rt),+) -> $t {
                $te { $($r: $r),+ }
            }
        }
        )+
    )
);
<commit_msg>Create the `u32_to_int!` and `usize_to_int!` macros<commit_after>macro_rules! impl_raw_accessors(
    ($(($t:ty, $raw:ty)),+) => (
        $(
        impl $t {
            #[inline]
            pub unsafe fn raw(&self) -> $raw { self.raw }
        }
        )+
    )
);

macro_rules! impl_owned_accessors(
    ($(($t:ty, $owned:ident)),+) => (
        $(
        impl $t {
            #[inline]
            pub unsafe fn $owned(&self) -> bool { self.$owned }
        }
        )+
    )
);

macro_rules! impl_raw_constructor(
    ($(($t:ty, $te:ident ($($r:ident:$rt:ty),+))),+) => (
        $(
        impl $t {
            #[inline]
            pub unsafe fn from_ll($($r:$rt),+) -> $t {
                $te { $($r: $r),+ }
            }
        }
        )+
    )
);

/// Many SDL functions will accept `int` values, even if it doesn't make sense for the values to be negative.
/// In the cases that SDL doesn't check negativity, passing negative values could be unsafe.
/// For example, `SDL_JoystickGetButton` uses the index argument to access an array without checking if it's negative,
/// which could potentially lead to segmentation faults.
macro_rules! u32_to_int(
    ($value:expr) => (
        if $value >= 1<<31 { Err(format!("`{}` is out of bounds.", stringify!($value))) }
        else { Ok($value as ::libc::c_int) }
    )
);

macro_rules! usize_to_int(
    ($value:expr) => (
        if $value >= 1<<31 { Err(format!("`{}` is out of bounds.", stringify!($value))) }
        else { Ok($value as ::libc::c_int) }
    )
);

macro_rules! int_to_u32(
    ($value:expr) => (
        if $value < 0 { Err(format!("`{}` is out of bounds.", stringify!($value))) }
        else { Ok($value as u32) }
    )
);
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::VecDeque;
use std::cmp;

use mio::{Token, EventLoop, EventSet, PollOpt, TryRead, TryWrite};
use mio::tcp::TcpStream;
use bytes::{Buf, MutBuf, ByteBuf, MutByteBuf, alloc};
use protobuf::Message as PbMessage;

use kvproto::msgpb::Message;
use kvproto::raft_serverpb::RaftSnapshotData;
use super::{Result, ConnData};
use super::server::Server;
use util::codec::rpc;
use super::transport::RaftStoreRouter;
use super::resolve::StoreAddrResolver;
use super::snap::Task as SnapTask;
use util::worker::Scheduler;

#[derive(PartialEq)]
enum ConnType {
    Handshake,
    Rpc,
    Snapshot,
}

const SNAPSHOT_PAYLOAD_BUF: usize = 4 * 1024 * 1024;

pub struct Conn {
    pub sock: TcpStream,
    pub token: Token,
    pub interest: EventSet,

    conn_type: ConnType,

    // store id is for remote store, we only set this
    // when we connect to the remote store.
    pub store_id: Option<u64>,

    // message header
    last_msg_id: u64,
    header: MutByteBuf,
    // message
    payload: Option<MutByteBuf>,

    file_size: usize,
    read_size: usize,
    snap_scheduler: Scheduler<SnapTask>,

    // write buffer, including msg header already.
    res: VecDeque<ByteBuf>,
}

fn try_read_data<T: TryRead, B: MutBuf>(r: &mut T, buf: &mut B) -> Result<()> {
    if buf.remaining() == 0 {
        return Ok(());
    }

    // TODO: use try_read_buf directly if we can solve the compile problem.
    unsafe {
        // header is not full read, we will try read more.
        if let Some(n) = try!(r.try_read(buf.mut_bytes())) {
            if n == 0 {
                // 0 means remote has closed the socket.
                return Err(box_err!("remote has closed the connection"));
            }
            buf.advance(n)
        }
    }

    Ok(())
}

fn create_mem_buf(s: usize) -> MutByteBuf {
    unsafe {
        ByteBuf::from_mem_ref(alloc::heap(s.next_power_of_two()), s as u32, 0, s as u32).flip()
    }
}

impl Conn {
    pub fn new(sock: TcpStream,
               token: Token,
               store_id: Option<u64>,
               snap_scheduler: Scheduler<SnapTask>)
               -> Conn {
        Conn {
            sock: sock,
            token: token,
            interest: EventSet::readable() | EventSet::hup(),
            conn_type: ConnType::Handshake,
            header: create_mem_buf(rpc::MSG_HEADER_LEN),
            read_size: 0,
            file_size: 0,
            payload: None,
            res: VecDeque::new(),
            last_msg_id: 0,
            snap_scheduler: snap_scheduler,
            store_id: store_id,
        }
    }

    pub fn close(&mut self) {
        if self.conn_type == ConnType::Snapshot {
            if let Err(e) = self.snap_scheduler.schedule(SnapTask::Discard(self.token)) {
                error!("failed to cleanup snapshot: {:?}", e);
            }
        }
    }

    pub fn reregister<T, S>(&mut self, event_loop: &mut EventLoop<Server<T, S>>) -> Result<()>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        try!(event_loop.reregister(&self.sock, self.token, self.interest, PollOpt::edge()));
        Ok(())
    }

    pub fn on_readable<T, S>(&mut self,
                             event_loop: &mut EventLoop<Server<T, S>>)
                             -> Result<Vec<ConnData>>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        let mut bufs = vec![];
        match self.conn_type {
            ConnType::Handshake => try!(self.handshake(event_loop, &mut bufs)),
            ConnType::Rpc => try!(self.read_rpc(event_loop, &mut bufs)),
            ConnType::Snapshot => try!(self.read_snapshot(event_loop)),
        };
        Ok(bufs)
    }

    fn handshake<T, S>(&mut self,
                       event_loop: &mut EventLoop<Server<T, S>>,
                       bufs: &mut Vec<ConnData>)
                       -> Result<()>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        let mut data = match try!(self.read_one_message()) {
            Some(data) => data,
            None => return Ok(()),
        };
        if data.is_snapshot() {
            self.conn_type = ConnType::Snapshot;

            let mut snap_data = RaftSnapshotData::new();
            try!(snap_data.merge_from_bytes(
                data.msg.get_raft().get_message().get_snapshot().get_data()));
            self.file_size = snap_data.get_file_size() as usize;
            self.payload = Some(create_mem_buf(cmp::min(SNAPSHOT_PAYLOAD_BUF, self.file_size)));

            let register_task = SnapTask::Register(self.token, data.msg.take_raft());
            box_try!(self.snap_scheduler.schedule(register_task));

            return self.read_snapshot(event_loop);
        }
        bufs.push(data);
        self.conn_type = ConnType::Rpc;
        self.read_rpc(event_loop, bufs)
    }

    fn read_snapshot<T, S>(&mut self, _: &mut EventLoop<Server<T, S>>) -> Result<()>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        // all content should be read, ignore any read operation.
        if self.payload.is_none() {
            return Ok(());
        }
        // TODO: limit rate
        while try!(self.read_payload()) {
            let payload = self.payload.take().unwrap();
            let cap = payload.capacity();
            self.read_size += cap;

            let task = SnapTask::Write(self.token, payload.flip());
            box_try!(self.snap_scheduler.schedule(task));

            if self.read_size == self.file_size {
                // last chunk
                box_try!(self.snap_scheduler.schedule(SnapTask::Close(self.token)));
                // let snap_scheduler to close the connection.
                break;
            } else if self.read_size + cap >= self.file_size {
                self.payload = Some(create_mem_buf(self.file_size - self.read_size))
            } else {
                self.payload = Some(create_mem_buf(cap))
            };
        }
        Ok(())
    }

    fn read_payload(&mut self) -> Result<bool> {
        let payload = self.payload.as_mut().unwrap();
        try!(try_read_data(&mut self.sock, payload));
        let ret = payload.remaining() == 0;
        Ok(ret)
    }

    fn read_one_message(&mut self) -> Result<Option<ConnData>> {
        if self.payload.is_none() {
            try!(try_read_data(&mut self.sock, &mut self.header));
            if self.header.remaining() > 0 {
                // we need to read more data for header
                return Ok(None);
            }

            // we have already read whole header, parse it and begin to read payload.
            let (msg_id, payload_len) = try!(rpc::decode_msg_header(self.header
                .bytes()));
            self.last_msg_id = msg_id;
            self.payload = Some(create_mem_buf(payload_len));
        }

        // payload here can't be None.
        let mut payload = self.payload.take().unwrap();
        try!(try_read_data(&mut self.sock, &mut payload));
        if payload.remaining() > 0 {
            // we need to read more data for payload
            self.payload = Some(payload);
            return Ok(None);
        }

        let mut msg = Message::new();
        try!(rpc::decode_body(payload.bytes(), &mut msg));
        self.header.clear();
        Ok(Some(ConnData {
            msg_id: self.last_msg_id,
            msg: msg,
        }))
    }

    fn read_rpc<T, S>(&mut self,
                      _: &mut EventLoop<Server<T, S>>,
                      bufs: &mut Vec<ConnData>)
                      -> Result<()>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        loop {
            // Because we use the edge trigger, so here we must read whole data.
            match try!(self.read_one_message()) {
                None => break,
                Some(d) => bufs.push(d),
            };
        }

        Ok(())
    }

    fn write_buf(&mut self) -> Result<usize> {
        // we check empty before.
        let mut buf = self.res.front_mut().unwrap();

        if let Some(n) = try!(self.sock.try_write(buf.bytes())) {
            buf.advance(n)
        }

        Ok(buf.remaining())
    }

    pub fn on_writable<T, S>(&mut self, event_loop: &mut EventLoop<Server<T, S>>) -> Result<()>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        while !self.res.is_empty() {
            let remaining = try!(self.write_buf());

            if remaining > 0 {
                // we don't write all data, so must try later.
                // we have already registered writable, no need registering again.
                return Ok(());
            }
            self.res.pop_front();
        }

        // no data for writing, remove writable
        self.interest.remove(EventSet::writable());
        try!(self.reregister(event_loop));

        Ok(())
    }

    pub fn append_write_buf<T, S>(&mut self,
                                  event_loop: &mut EventLoop<Server<T, S>>,
                                  msg: ConnData)
                                  -> Result<()>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        // Now we just push data to a write buffer and register writable for later writing.
        // Later we can write data directly, if meet WOUNDBLOCK error(don't write all data OK),
        // we can register writable at that time.
        // We must also check `socket is not connected` error too, when we connect to a remote
        // store, mio puts this socket in event loop immediately, but this socket may not be
        // connected at that time, so we must register writable too for this case.
        self.res.push_back(msg.encode_to_buf());

        if !self.interest.is_writable() {
            // re-register writable if we have not,
            // if registered, we can only remove this flag when
            // writing all data in writable function.
            self.interest.insert(EventSet::writable());
            try!(self.reregister(event_loop));
        }

        Ok(())
    }
}
<commit_msg>server: cleanup try_read/write and create_mem_buf. (#760)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::VecDeque;
use std::cmp;

use mio::{Token, EventLoop, EventSet, PollOpt};
use mio::tcp::TcpStream;
use bytes::{Buf, MutBuf, ByteBuf, MutByteBuf};
use protobuf::Message as PbMessage;

use kvproto::msgpb::Message;
use kvproto::raft_serverpb::RaftSnapshotData;
use super::{Result, ConnData};
use super::server::Server;
use util::codec::rpc;
use super::transport::RaftStoreRouter;
use super::resolve::StoreAddrResolver;
use super::snap::Task as SnapTask;
use util::worker::Scheduler;
use util::buf::{TryRead, TryWrite, create_mem_buf};

#[derive(PartialEq)]
enum ConnType {
    Handshake,
    Rpc,
    Snapshot,
}

const SNAPSHOT_PAYLOAD_BUF: usize = 4 * 1024 * 1024;

pub struct Conn {
    pub sock: TcpStream,
    pub token: Token,
    pub interest: EventSet,

    conn_type: ConnType,

    // store id is for remote store, we only set this
    // when we connect to the remote store.
    pub store_id: Option<u64>,

    // message header
    last_msg_id: u64,
    header: MutByteBuf,
    // message
    payload: Option<MutByteBuf>,

    file_size: usize,
    read_size: usize,
    snap_scheduler: Scheduler<SnapTask>,

    // write buffer, including msg header already.
    res: VecDeque<ByteBuf>,
}

fn try_read_data<T: TryRead, B: MutBuf>(r: &mut T, buf: &mut B) -> Result<()> {
    if buf.remaining() == 0 {
        return Ok(());
    }

    if let Some(n) = try!(r.try_read_buf(buf)) {
        if n == 0 {
            // 0 means remote has closed the socket.
            return Err(box_err!("remote has closed the connection"));
        }
    }

    Ok(())
}

impl Conn {
    pub fn new(sock: TcpStream,
               token: Token,
               store_id: Option<u64>,
               snap_scheduler: Scheduler<SnapTask>)
               -> Conn {
        Conn {
            sock: sock,
            token: token,
            interest: EventSet::readable() | EventSet::hup(),
            conn_type: ConnType::Handshake,
            header: create_mem_buf(rpc::MSG_HEADER_LEN),
            read_size: 0,
            file_size: 0,
            payload: None,
            res: VecDeque::new(),
            last_msg_id: 0,
            snap_scheduler: snap_scheduler,
            store_id: store_id,
        }
    }

    pub fn close(&mut self) {
        if self.conn_type == ConnType::Snapshot {
            if let Err(e) = self.snap_scheduler.schedule(SnapTask::Discard(self.token)) {
                error!("failed to cleanup snapshot: {:?}", e);
            }
        }
    }

    pub fn reregister<T, S>(&mut self, event_loop: &mut EventLoop<Server<T, S>>) -> Result<()>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        try!(event_loop.reregister(&self.sock, self.token, self.interest, PollOpt::edge()));
        Ok(())
    }

    pub fn on_readable<T, S>(&mut self,
                             event_loop: &mut EventLoop<Server<T, S>>)
                             -> Result<Vec<ConnData>>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        let mut bufs = vec![];
        match self.conn_type {
            ConnType::Handshake => try!(self.handshake(event_loop, &mut bufs)),
            ConnType::Rpc => try!(self.read_rpc(event_loop, &mut bufs)),
            ConnType::Snapshot => try!(self.read_snapshot(event_loop)),
        };
        Ok(bufs)
    }

    fn handshake<T, S>(&mut self,
                       event_loop: &mut EventLoop<Server<T, S>>,
                       bufs: &mut Vec<ConnData>)
                       -> Result<()>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        let mut data = match try!(self.read_one_message()) {
            Some(data) => data,
            None => return Ok(()),
        };
        if data.is_snapshot() {
            self.conn_type = ConnType::Snapshot;

            let mut snap_data = RaftSnapshotData::new();
            try!(snap_data.merge_from_bytes(
                data.msg.get_raft().get_message().get_snapshot().get_data()));
            self.file_size = snap_data.get_file_size() as usize;
            self.payload = Some(create_mem_buf(cmp::min(SNAPSHOT_PAYLOAD_BUF, self.file_size)));

            let register_task = SnapTask::Register(self.token, data.msg.take_raft());
            box_try!(self.snap_scheduler.schedule(register_task));

            return self.read_snapshot(event_loop);
        }
        bufs.push(data);
        self.conn_type = ConnType::Rpc;
        self.read_rpc(event_loop, bufs)
    }

    fn read_snapshot<T, S>(&mut self, _: &mut EventLoop<Server<T, S>>) -> Result<()>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        // all content should be read, ignore any read operation.
        if self.payload.is_none() {
            return Ok(());
        }
        // TODO: limit rate
        while try!(self.read_payload()) {
            let payload = self.payload.take().unwrap();
            let cap = payload.capacity();
            self.read_size += cap;

            let task = SnapTask::Write(self.token, payload.flip());
            box_try!(self.snap_scheduler.schedule(task));

            if self.read_size == self.file_size {
                // last chunk
                box_try!(self.snap_scheduler.schedule(SnapTask::Close(self.token)));
                // let snap_scheduler to close the connection.
                break;
            } else if self.read_size + cap >= self.file_size {
                self.payload = Some(create_mem_buf(self.file_size - self.read_size))
            } else {
                self.payload = Some(create_mem_buf(cap))
            };
        }
        Ok(())
    }

    fn read_payload(&mut self) -> Result<bool> {
        let payload = self.payload.as_mut().unwrap();
        try!(try_read_data(&mut self.sock, payload));
        let ret = payload.remaining() == 0;
        Ok(ret)
    }

    fn read_one_message(&mut self) -> Result<Option<ConnData>> {
        if self.payload.is_none() {
            try!(try_read_data(&mut self.sock, &mut self.header));
            if self.header.remaining() > 0 {
                // we need to read more data for header
                return Ok(None);
            }

            // we have already read whole header, parse it and begin to read payload.
            let (msg_id, payload_len) = try!(rpc::decode_msg_header(self.header
                .bytes()));
            self.last_msg_id = msg_id;
            self.payload = Some(create_mem_buf(payload_len));
        }

        // payload here can't be None.
        let mut payload = self.payload.take().unwrap();
        try!(try_read_data(&mut self.sock, &mut payload));
        if payload.remaining() > 0 {
            // we need to read more data for payload
            self.payload = Some(payload);
            return Ok(None);
        }

        let mut msg = Message::new();
        try!(rpc::decode_body(payload.bytes(), &mut msg));
        self.header.clear();
        Ok(Some(ConnData {
            msg_id: self.last_msg_id,
            msg: msg,
        }))
    }

    fn read_rpc<T, S>(&mut self,
                      _: &mut EventLoop<Server<T, S>>,
                      bufs: &mut Vec<ConnData>)
                      -> Result<()>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        loop {
            // Because we use the edge trigger, so here we must read whole data.
            match try!(self.read_one_message()) {
                None => break,
                Some(d) => bufs.push(d),
            };
        }

        Ok(())
    }

    fn write_buf(&mut self) -> Result<usize> {
        // we check empty before.
        let mut buf = self.res.front_mut().unwrap();

        try!(self.sock.try_write_buf(buf));

        Ok(buf.remaining())
    }

    pub fn on_writable<T, S>(&mut self, event_loop: &mut EventLoop<Server<T, S>>) -> Result<()>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        while !self.res.is_empty() {
            let remaining = try!(self.write_buf());

            if remaining > 0 {
                // we don't write all data, so must try later.
                // we have already registered writable, no need registering again.
                return Ok(());
            }
            self.res.pop_front();
        }

        // no data for writing, remove writable
        self.interest.remove(EventSet::writable());
        try!(self.reregister(event_loop));

        Ok(())
    }

    pub fn append_write_buf<T, S>(&mut self,
                                  event_loop: &mut EventLoop<Server<T, S>>,
                                  msg: ConnData)
                                  -> Result<()>
        where T: RaftStoreRouter,
              S: StoreAddrResolver
    {
        // Now we just push data to a write buffer and register writable for later writing.
        // Later we can write data directly, if meet WOUNDBLOCK error(don't write all data OK),
        // we can register writable at that time.
        // We must also check `socket is not connected` error too, when we connect to a remote
        // store, mio puts this socket in event loop immediately, but this socket may not be
        // connected at that time, so we must register writable too for this case.
        self.res.push_back(msg.encode_to_buf());

        if !self.interest.is_writable() {
            // re-register writable if we have not,
            // if registered, we can only remove this flag when
            // writing all data in writable function.
            self.interest.insert(EventSet::writable());
            try!(self.reregister(event_loop));
        }

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


<commit_before>mod error;

use std::convert::{TryFrom, TryInto};

use fleetspeak::Packet;
use log::error;

use crate::action;
pub use self::error::{Error, ParseError};

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

pub fn execute<S, R, H>(session: &mut S, handler: H, payload: Payload) -> Result<()>
where
    S: Session,
    R: action::Request,
    H: FnOnce(&mut S, R) -> Result<()>,
{
    handler(session, payload.parse()?)
}

pub fn handle<M>(message: M)
where
    M: TryInto<Demand, Error=ParseError>,
{
    let demand = match message.try_into() {
        Ok(demand) => demand,
        Err(error) => {
            error!("failed to parse the message: {}", error);
            return;
        }
    };

    let mut session = Action::new(demand.header.clone());
    let result = action::dispatch(&demand.action, &mut session, demand.payload);

    let status = Status {
        header: demand.header,
        result: result,
    };

    let message = match status.try_into() {
        Ok(message) => message,
        Err(error) => {
            // If we cannot encode the final status message, there is nothing
            // we can do to notify the server, as status is responsible for
            // reporting errors. We can only log the error and carry on.
            error!("failed to encode status message: {}", error);
            return;
        }
    };

    send(message);
}

pub trait Session {
    fn send<R: action::Response>(&mut self, response: R) -> Result<()>;
}

pub struct Action {
    header: Header,
    next_response_id: u64,
}

impl Action {

    pub fn new(header: Header) -> Action {
        Action {
            header: header,
            next_response_id: 0,
        }
    }
}

impl Session for Action {

    fn send<R: action::Response>(&mut self, response: R) -> Result<()> {
        Response {
            session_id: self.header.session_id.clone(),
            request_id: Some(self.header.request_id),
            response_id: Some(self.next_response_id),
            data: response,
        }.send()?;

        self.next_response_id += 1;

        Ok(())
    }
}

pub struct Sink {
    id: &'static str,
}

impl Session for Sink {

    fn send<R: action::Response>(&mut self, response: R) -> Result<()> {
        Response {
            session_id: String::from(self.id),
            request_id: None,
            response_id: None,
            data: response,
        }.send()
    }
}

pub fn startup() -> Sink {
    Sink { id: "/flows/F:Startup" }
}

pub struct Demand {
    pub action: String,
    pub header: Header,
    pub payload: Payload,
}

impl TryFrom<rrg_proto::GrrMessage> for Demand {

    type Error = ParseError;

    fn try_from(message: rrg_proto::GrrMessage)
    -> std::result::Result<Demand, ParseError>
    {
        use ParseError::*;

        let header = Header {
            session_id: message.session_id.ok_or(MissingField("session id"))?,
            request_id: message.request_id.ok_or(MissingField("request id"))?,
        };

        Ok(Demand {
            action: message.name.ok_or(MissingField("action name"))?,
            header: header,
            payload: Payload(message.args),
        })
    }
}

#[derive(Clone, Debug)]
pub struct Header {
    pub session_id: String,
    pub request_id: u64,
}

#[derive(Debug)]
pub struct Payload(Option<Vec<u8>>);

impl Payload {

    pub fn parse<R>(&self) -> std::result::Result<R, ParseError>
    where
        R: action::Request,
    {
        let proto = match self {
            Payload(Some(bytes)) => prost::Message::decode(&bytes[..])?,
            Payload(None) => Default::default(),
        };

        Ok(R::from_proto(proto))
    }
}

struct Response<R: action::Response> {
    session_id: String,
    request_id: Option<u64>,
    response_id: Option<u64>,
    data: R,
}

impl<R: action::Response> Response<R> {

    fn send(self) -> Result<()> {
        let message = self.try_into()?;
        send(message);

        Ok(())
    }
}

impl<R: action::Response> TryInto<rrg_proto::GrrMessage> for Response<R> {

    type Error = prost::EncodeError;

    fn try_into(self)
    -> std::result::Result<rrg_proto::GrrMessage, prost::EncodeError>
    {
        let mut data = Vec::new();
        prost::Message::encode(&self.data.into_proto(), &mut data)?;

        Ok(rrg_proto::GrrMessage {
            session_id: Some(self.session_id),
            response_id: self.response_id,
            request_id: self.request_id,
            r#type: Some(rrg_proto::grr_message::Type::Message.into()),
            args_rdf_name: R::RDF_NAME.map(String::from),
            args: Some(data),
            ..Default::default()
        })
    }
}

struct Status {
    header: Header,
    result: Result<()>,
}

impl TryInto<rrg_proto::GrrMessage> for Status {

    type Error = prost::EncodeError;

    fn try_into(self)
    -> std::result::Result<rrg_proto::GrrMessage, prost::EncodeError> {
        let status = match self.result {
            Ok(()) => rrg_proto::GrrStatus {
                status: Some(rrg_proto::grr_status::ReturnedStatus::Ok.into()),
                ..Default::default()
            },
            Err(error) => rrg_proto::GrrStatus {
                status: Some(rrg_proto::grr_status::ReturnedStatus::GenericError.into()),
                error_message: Some(error.to_string()),
                ..Default::default()
            },
        };

        let mut data = Vec::new();
        prost::Message::encode(&status, &mut data)?;

        Ok(rrg_proto::GrrMessage {
            session_id: Some(self.header.session_id),
            response_id: Some(self.header.request_id),
            r#type: Some(rrg_proto::grr_message::Type::Status.into()),
            args_rdf_name: Some(String::from("GrrStatus")),
            args: Some(data),
            ..Default::default()
        })
    }
}

fn send(message: rrg_proto::GrrMessage) {
        let packet = Packet {
            service: String::from("GRR"),
            kind: Some(String::from("GrrMessage")),
            data: message,
        };

        if let Err(error) = fleetspeak::send(packet) {
            // If we failed to deliver the message through Fleetspeak, it means
            // that our communication is broken (e.g. the pipe was closed) and
            // the agent should be killed.
            panic!("message delivery failure: {}", error)
        };
}
<commit_msg>Add missing license header.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

mod error;

use std::convert::{TryFrom, TryInto};

use fleetspeak::Packet;
use log::error;

use crate::action;
pub use self::error::{Error, ParseError};

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

pub fn execute<S, R, H>(session: &mut S, handler: H, payload: Payload) -> Result<()>
where
    S: Session,
    R: action::Request,
    H: FnOnce(&mut S, R) -> Result<()>,
{
    handler(session, payload.parse()?)
}

pub fn handle<M>(message: M)
where
    M: TryInto<Demand, Error=ParseError>,
{
    let demand = match message.try_into() {
        Ok(demand) => demand,
        Err(error) => {
            error!("failed to parse the message: {}", error);
            return;
        }
    };

    let mut session = Action::new(demand.header.clone());
    let result = action::dispatch(&demand.action, &mut session, demand.payload);

    let status = Status {
        header: demand.header,
        result: result,
    };

    let message = match status.try_into() {
        Ok(message) => message,
        Err(error) => {
            // If we cannot encode the final status message, there is nothing
            // we can do to notify the server, as status is responsible for
            // reporting errors. We can only log the error and carry on.
            error!("failed to encode status message: {}", error);
            return;
        }
    };

    send(message);
}

pub trait Session {
    fn send<R: action::Response>(&mut self, response: R) -> Result<()>;
}

pub struct Action {
    header: Header,
    next_response_id: u64,
}

impl Action {

    pub fn new(header: Header) -> Action {
        Action {
            header: header,
            next_response_id: 0,
        }
    }
}

impl Session for Action {

    fn send<R: action::Response>(&mut self, response: R) -> Result<()> {
        Response {
            session_id: self.header.session_id.clone(),
            request_id: Some(self.header.request_id),
            response_id: Some(self.next_response_id),
            data: response,
        }.send()?;

        self.next_response_id += 1;

        Ok(())
    }
}

pub struct Sink {
    id: &'static str,
}

impl Session for Sink {

    fn send<R: action::Response>(&mut self, response: R) -> Result<()> {
        Response {
            session_id: String::from(self.id),
            request_id: None,
            response_id: None,
            data: response,
        }.send()
    }
}

pub fn startup() -> Sink {
    Sink { id: "/flows/F:Startup" }
}

pub struct Demand {
    pub action: String,
    pub header: Header,
    pub payload: Payload,
}

impl TryFrom<rrg_proto::GrrMessage> for Demand {

    type Error = ParseError;

    fn try_from(message: rrg_proto::GrrMessage)
    -> std::result::Result<Demand, ParseError>
    {
        use ParseError::*;

        let header = Header {
            session_id: message.session_id.ok_or(MissingField("session id"))?,
            request_id: message.request_id.ok_or(MissingField("request id"))?,
        };

        Ok(Demand {
            action: message.name.ok_or(MissingField("action name"))?,
            header: header,
            payload: Payload(message.args),
        })
    }
}

#[derive(Clone, Debug)]
pub struct Header {
    pub session_id: String,
    pub request_id: u64,
}

#[derive(Debug)]
pub struct Payload(Option<Vec<u8>>);

impl Payload {

    pub fn parse<R>(&self) -> std::result::Result<R, ParseError>
    where
        R: action::Request,
    {
        let proto = match self {
            Payload(Some(bytes)) => prost::Message::decode(&bytes[..])?,
            Payload(None) => Default::default(),
        };

        Ok(R::from_proto(proto))
    }
}

struct Response<R: action::Response> {
    session_id: String,
    request_id: Option<u64>,
    response_id: Option<u64>,
    data: R,
}

impl<R: action::Response> Response<R> {

    fn send(self) -> Result<()> {
        let message = self.try_into()?;
        send(message);

        Ok(())
    }
}

impl<R: action::Response> TryInto<rrg_proto::GrrMessage> for Response<R> {

    type Error = prost::EncodeError;

    fn try_into(self)
    -> std::result::Result<rrg_proto::GrrMessage, prost::EncodeError>
    {
        let mut data = Vec::new();
        prost::Message::encode(&self.data.into_proto(), &mut data)?;

        Ok(rrg_proto::GrrMessage {
            session_id: Some(self.session_id),
            response_id: self.response_id,
            request_id: self.request_id,
            r#type: Some(rrg_proto::grr_message::Type::Message.into()),
            args_rdf_name: R::RDF_NAME.map(String::from),
            args: Some(data),
            ..Default::default()
        })
    }
}

struct Status {
    header: Header,
    result: Result<()>,
}

impl TryInto<rrg_proto::GrrMessage> for Status {

    type Error = prost::EncodeError;

    fn try_into(self)
    -> std::result::Result<rrg_proto::GrrMessage, prost::EncodeError> {
        let status = match self.result {
            Ok(()) => rrg_proto::GrrStatus {
                status: Some(rrg_proto::grr_status::ReturnedStatus::Ok.into()),
                ..Default::default()
            },
            Err(error) => rrg_proto::GrrStatus {
                status: Some(rrg_proto::grr_status::ReturnedStatus::GenericError.into()),
                error_message: Some(error.to_string()),
                ..Default::default()
            },
        };

        let mut data = Vec::new();
        prost::Message::encode(&status, &mut data)?;

        Ok(rrg_proto::GrrMessage {
            session_id: Some(self.header.session_id),
            response_id: Some(self.header.request_id),
            r#type: Some(rrg_proto::grr_message::Type::Status.into()),
            args_rdf_name: Some(String::from("GrrStatus")),
            args: Some(data),
            ..Default::default()
        })
    }
}

fn send(message: rrg_proto::GrrMessage) {
        let packet = Packet {
            service: String::from("GRR"),
            kind: Some(String::from("GrrMessage")),
            data: message,
        };

        if let Err(error) = fleetspeak::send(packet) {
            // If we failed to deliver the message through Fleetspeak, it means
            // that our communication is broken (e.g. the pipe was closed) and
            // the agent should be killed.
            panic!("message delivery failure: {}", error)
        };
}
<|endoftext|>