<commit_before>//! JIT Compiler to generate code fragments in runtime.

extern crate libc;
extern crate llvm_sys;

#[macro_use]
mod macros;
mod block;
mod buffer;
mod builder;
mod util;
mod types;
mod value;

use std::mem;
use std::ptr;

use llvm_sys::core;
use llvm_sys::prelude::{
  LLVMContextRef,
  LLVMModuleRef
};
use llvm_sys::bit_reader::LLVMParseBitcodeInContext;
use llvm_sys::execution_engine::{
  LLVMExecutionEngineRef,
  LLVMLinkInMCJIT,
  LLVMMCJITCompilerOptions,
  LLVMCreateMCJITCompilerForModule
};
use llvm_sys::linker;
use llvm_sys::target::{
  LLVM_InitializeNativeTarget,
  LLVM_InitializeNativeAsmPrinter
};
use llvm_sys::target_machine::LLVMCodeModel;

use libc::{c_char, c_uint};

use buffer::MemoryBuffer;
use builder::Builder;
use value::{GlobalValue, Value, ValueIter};
use types::Ty;
use util::chars;

pub const JIT_OPT_LVEL: usize = 2;

#[repr(C)]
#[derive(Copy, Clone)]
pub enum AddressSpace
{
  Generic = 0,
  Global = 1,
  Shared = 3,
  Const = 4,
  Local = 5,
}

fn new_module(ctx: LLVMContextRef, name: &str) -> LLVMModuleRef
{
  let c_name = util::chars::from_str(name);
  unsafe {
  	core::LLVMModuleCreateWithNameInContext(c_name, ctx)
  }
}
fn new_module_from_bc(ctx: LLVMContextRef, path: &str) -> Result<LLVMModuleRef, String>
{
  unsafe {
    let mut out: LLVMModuleRef = mem::uninitialized();
    let mut err: *mut c_char   = mem::uninitialized();
    let buf = try!(MemoryBuffer::from_file(path));

    let ret = LLVMParseBitcodeInContext(ctx,
                                        buf.as_ptr(),
    	                                  &mut out,
    	                                  &mut err);
    llvm_ret!(ret, out, err)
  }
}

fn new_jit_ee(m: LLVMModuleRef, opt_lv: usize) -> Result<LLVMExecutionEngineRef, String>
{
  unsafe {
    let mut ee : LLVMExecutionEngineRef = mem::uninitialized();
    let mut err: *mut c_char = mem::uninitialized();

    LLVMLinkInMCJIT();
    expect_noerr!(LLVM_InitializeNativeTarget(), "failed to initialize native target");
    expect_noerr!(LLVM_InitializeNativeAsmPrinter(), "failed to initialize native asm printer");

    let mut opts = new_mcjit_compiler_options(opt_lv);
    let opts_size = mem::size_of::<LLVMMCJITCompilerOptions>();

    let ret = LLVMCreateMCJITCompilerForModule(&mut ee,
      		                                     m,
      		                                     &mut opts,
      		                                     opts_size as u64,
      		                                     &mut err);
    llvm_ret!(ret, ee, err)
  }
}

fn new_mcjit_compiler_options(opt_lv: usize) -> LLVMMCJITCompilerOptions
{
  LLVMMCJITCompilerOptions {
    OptLevel: opt_lv as c_uint,
    CodeModel: LLVMCodeModel::LLVMCodeModelJITDefault,
    NoFramePointerElim: 0,
    EnableFastISel: 1,
    MCJMM: ptr::null_mut()
  }
}

pub struct JitCompiler
{
  ctx    : LLVMContextRef,
  module : LLVMModuleRef,
  ee     : LLVMExecutionEngineRef,
  //builder: Builder // internal use
}

impl JitCompiler {

  pub fn new(module_name: &str) -> Result<JitCompiler, String>
  {
    let ctx     = unsafe { core::LLVMContextCreate() };
    let module  = new_module(ctx, module_name);
    JitCompiler::new_internal(ctx, module)
  }

  pub fn new_from_bc(bitcode_path: &str) -> Result<JitCompiler, String>
  {
    let ctx     = unsafe { core::LLVMContextCreate() };
    let module  = try!(new_module_from_bc(ctx, bitcode_path));
    JitCompiler::new_internal(ctx, module)
  }

  fn new_internal(ctx: LLVMContextRef, module: LLVMModuleRef) -> Result<JitCompiler, String>
  {
    let ee      = try!(new_jit_ee(module, JIT_OPT_LVEL));
    //let builder = Builder(unsafe { core::LLVMCreateBuilderInContext(ctx) });

    Ok(JitCompiler {
      ctx    : ctx,
      module : module,
      ee     : ee,
      //builder: builder
    })
  }

  pub fn context(&self) -> LLVMContextRef { self.ctx }
  pub fn module(&self) -> LLVMModuleRef { self.module }
  pub fn engine(&self) -> LLVMExecutionEngineRef { self.ee }

  /// Add an external global to the module with the given type and name.
  pub fn add_global(&self, name: &str, ty: &Ty) -> GlobalValue
  {
    let c_name = chars::from_str(name);
    GlobalValue(
      unsafe {
        core::LLVMAddGlobal(self.module, ty.0, c_name)
      }
    )
  }

  /// Add a global in the given address space to the module with the given type and name.
  pub fn add_global_in_addr_space(&self,
  	                              name: &str,
  	                              ty: &Ty,
  	                              sp: AddressSpace) -> GlobalValue
  {
    let c_name = chars::from_str(name);
    GlobalValue(
      unsafe {
        core::LLVMAddGlobalInAddressSpace(self.module, ty.0, c_name, sp as c_uint)
      }
    )
  }

  /// Add a constant global to the module with the given type, name and value.
  pub fn add_global_constant(&self, name: &str, val: &Value) -> GlobalValue
  {
    let c_name = chars::from_str(name);
    GlobalValue(
      unsafe {
        let global = core::LLVMAddGlobal(self.module, val.ty().0, c_name);
        core::LLVMSetInitializer (global, val.0);
        global
      }
    )
  }

  /// Get the global with the name given, or `None` if no global with that name exists.
  pub fn get_global(&self, name: &str) -> Option<GlobalValue>
  {
    let c_name = chars::from_str(name);
    unsafe {
      let ptr = core::LLVMGetNamedGlobal(self.module, c_name);
      if ptr.is_null() {
        None
      } else {
        Some(GlobalValue(ptr))
      }
    }
  }

  /// Get an iterator of global values
  pub fn global_values(&self) -> ValueIter<GlobalValue>
  {
  	ValueIter::new(
 			unsafe { core::LLVMGetFirstGlobal(self.module) },
 			core::LLVMGetNextGlobal
 		)
  }

  /// Link a module into this module, returning an error string if an error occurs.
  ///
  /// This *does not* destroy the source module.
  pub fn link(&self, module: LLVMModuleRef) -> Result<(), String>
  {
    unsafe {
      let mut error = mem::uninitialized();
      let ret = linker::LLVMLinkModules(self.module, module,
      	                                linker::LLVMLinkerMode::LLVMLinkerPreserveSource,
      	                                &mut error);
      llvm_ret!(ret, (), error)
    }
  }

  /// Link a module into this module, returning an error string if an error occurs.
  ///
  /// This *does* destroy the source module.
  pub fn link_destroy(&self, module: LLVMModuleRef) -> Result<(), String>
  {
    unsafe {
      let mut error = mem::uninitialized();
      let ret = linker::LLVMLinkModules(self.module, module,
      	                                linker::LLVMLinkerMode::LLVMLinkerDestroySource,
      	                                &mut error);
      llvm_ret!(ret, (), error)
    }
  }

  pub fn new_builder(&self) -> Builder
  {
    Builder(unsafe { core::LLVMCreateBuilderInContext(self.ctx) })
  }
}

impl Drop for JitCompiler {
  fn drop(&mut self) {
    unsafe {
      core::LLVMDisposeModule(self.module);
      core::LLVMContextDispose(self.ctx);
    }
  }
}

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

  #[test]
  fn test_ctor() {
    let jit = JitCompiler::new("target/test-ir/test-module.bc").ok().unwrap();
  }
}<commit_msg>Add optimize() to JitCompiler.<commit_after>//! JIT Compiler to generate code fragments in runtime.

extern crate libc;
extern crate llvm_sys;

#[macro_use]
mod macros;
mod block;
mod buffer;
mod builder;
mod util;
mod types;
mod value;

use std::mem;
use std::ptr;

use llvm_sys::core;
use llvm_sys::prelude::{
  LLVMContextRef,
  LLVMModuleRef
};
use llvm_sys::bit_reader::LLVMParseBitcodeInContext;
use llvm_sys::execution_engine::{
  LLVMExecutionEngineRef,
  LLVMLinkInMCJIT,
  LLVMMCJITCompilerOptions,
  LLVMCreateMCJITCompilerForModule
};
use llvm_sys::linker;
use llvm_sys::transforms::pass_manager_builder as pass;
use llvm_sys::target::{
  LLVM_InitializeNativeTarget,
  LLVM_InitializeNativeAsmPrinter
};
use llvm_sys::target_machine::LLVMCodeModel;

use libc::{c_char, c_uint};

use buffer::MemoryBuffer;
use builder::Builder;
use value::{GlobalValue, Value, ValueIter};
use types::Ty;
use util::chars;

pub const JIT_OPT_LVEL: usize = 2;

#[repr(C)]
#[derive(Copy, Clone)]
pub enum AddressSpace
{
  Generic = 0,
  Global = 1,
  Shared = 3,
  Const = 4,
  Local = 5,
}

fn new_module(ctx: LLVMContextRef, name: &str) -> LLVMModuleRef
{
  let c_name = util::chars::from_str(name);
  unsafe {
  	core::LLVMModuleCreateWithNameInContext(c_name, ctx)
  }
}
fn new_module_from_bc(ctx: LLVMContextRef, path: &str) -> Result<LLVMModuleRef, String>
{
  unsafe {
    let mut out: LLVMModuleRef = mem::uninitialized();
    let mut err: *mut c_char   = mem::uninitialized();
    let buf = try!(MemoryBuffer::from_file(path));

    let ret = LLVMParseBitcodeInContext(ctx,
                                        buf.as_ptr(),
    	                                  &mut out,
    	                                  &mut err);
    llvm_ret!(ret, out, err)
  }
}

fn new_jit_ee(m: LLVMModuleRef, opt_lv: usize) -> Result<LLVMExecutionEngineRef, String>
{
  unsafe {
    let mut ee : LLVMExecutionEngineRef = mem::uninitialized();
    let mut err: *mut c_char = mem::uninitialized();

    LLVMLinkInMCJIT();
    expect_noerr!(LLVM_InitializeNativeTarget(), "failed to initialize native target");
    expect_noerr!(LLVM_InitializeNativeAsmPrinter(), "failed to initialize native asm printer");

    let mut opts = new_mcjit_compiler_options(opt_lv);
    let opts_size = mem::size_of::<LLVMMCJITCompilerOptions>();

    let ret = LLVMCreateMCJITCompilerForModule(&mut ee,
      		                                     m,
      		                                     &mut opts,
      		                                     opts_size as u64,
      		                                     &mut err);
    llvm_ret!(ret, ee, err)
  }
}

fn new_mcjit_compiler_options(opt_lv: usize) -> LLVMMCJITCompilerOptions
{
  LLVMMCJITCompilerOptions {
    OptLevel: opt_lv as c_uint,
    CodeModel: LLVMCodeModel::LLVMCodeModelJITDefault,
    NoFramePointerElim: 0,
    EnableFastISel: 1,
    MCJMM: ptr::null_mut()
  }
}

pub struct JitCompiler
{
  ctx    : LLVMContextRef,
  module : LLVMModuleRef,
  ee     : LLVMExecutionEngineRef,
  //builder: Builder // internal use
}

impl JitCompiler {

  pub fn new(module_name: &str) -> Result<JitCompiler, String>
  {
    let ctx     = unsafe { core::LLVMContextCreate() };
    let module  = new_module(ctx, module_name);
    JitCompiler::new_internal(ctx, module)
  }

  pub fn new_from_bc(bitcode_path: &str) -> Result<JitCompiler, String>
  {
    let ctx     = unsafe { core::LLVMContextCreate() };
    let module  = try!(new_module_from_bc(ctx, bitcode_path));
    JitCompiler::new_internal(ctx, module)
  }

  fn new_internal(ctx: LLVMContextRef, module: LLVMModuleRef) -> Result<JitCompiler, String>
  {
    let ee      = try!(new_jit_ee(module, JIT_OPT_LVEL));
    //let builder = Builder(unsafe { core::LLVMCreateBuilderInContext(ctx) });

    Ok(JitCompiler {
      ctx    : ctx,
      module : module,
      ee     : ee,
      //builder: builder
    })
  }

  pub fn context(&self) -> LLVMContextRef { self.ctx }
  pub fn module(&self) -> LLVMModuleRef { self.module }
  pub fn engine(&self) -> LLVMExecutionEngineRef { self.ee }

  /// Add an external global to the module with the given type and name.
  pub fn add_global(&self, name: &str, ty: &Ty) -> GlobalValue
  {
    let c_name = chars::from_str(name);
    GlobalValue(
      unsafe {
        core::LLVMAddGlobal(self.module, ty.0, c_name)
      }
    )
  }

  /// Add a global in the given address space to the module with the given type and name.
  pub fn add_global_in_addr_space(&self,
  	                              name: &str,
  	                              ty: &Ty,
  	                              sp: AddressSpace) -> GlobalValue
  {
    let c_name = chars::from_str(name);
    GlobalValue(
      unsafe {
        core::LLVMAddGlobalInAddressSpace(self.module, ty.0, c_name, sp as c_uint)
      }
    )
  }

  /// Add a constant global to the module with the given type, name and value.
  pub fn add_global_constant(&self, name: &str, val: &Value) -> GlobalValue
  {
    let c_name = chars::from_str(name);
    GlobalValue(
      unsafe {
        let global = core::LLVMAddGlobal(self.module, val.ty().0, c_name);
        core::LLVMSetInitializer (global, val.0);
        global
      }
    )
  }

  /// Get the global with the name given, or `None` if no global with that name exists.
  pub fn get_global(&self, name: &str) -> Option<GlobalValue>
  {
    let c_name = chars::from_str(name);
    unsafe {
      let ptr = core::LLVMGetNamedGlobal(self.module, c_name);
      if ptr.is_null() {
        None
      } else {
        Some(GlobalValue(ptr))
      }
    }
  }

  /// Get an iterator of global values
  pub fn global_values(&self) -> ValueIter<GlobalValue>
  {
  	ValueIter::new(
 			unsafe { core::LLVMGetFirstGlobal(self.module) },
 			core::LLVMGetNextGlobal
 		)
  }

  /// Link a module into this module, returning an error string if an error occurs.
  ///
  /// This *does not* destroy the source module.
  pub fn link(&self, module: LLVMModuleRef) -> Result<(), String>
  {
    unsafe {
      let mut error = mem::uninitialized();
      let ret = linker::LLVMLinkModules(self.module, module,
      	                                linker::LLVMLinkerMode::LLVMLinkerPreserveSource,
      	                                &mut error);
      llvm_ret!(ret, (), error)
    }
  }

  /// Link a module into this module, returning an error string if an error occurs.
  ///
  /// This *does* destroy the source module.
  pub fn link_destroy(&self, module: LLVMModuleRef) -> Result<(), String>
  {
    unsafe {
      let mut error = mem::uninitialized();
      let ret = linker::LLVMLinkModules(self.module, module,
      	                                linker::LLVMLinkerMode::LLVMLinkerDestroySource,
      	                                &mut error);
      llvm_ret!(ret, (), error)
    }
  }

  /// Optimize this module with the given optimization level and size level.
  ///
  /// This runs passes depending on the levels given.
  pub fn optimize(&self, opt_level: usize, size_level: usize) {
    unsafe {
      let builder = pass::LLVMPassManagerBuilderCreate();
      pass::LLVMPassManagerBuilderSetOptLevel(builder, opt_level as c_uint);
      pass::LLVMPassManagerBuilderSetSizeLevel(builder, size_level as c_uint);
      let pass_manager = core::LLVMCreatePassManager();
      pass::LLVMPassManagerBuilderPopulateModulePassManager(builder, pass_manager);
      pass::LLVMPassManagerBuilderDispose(builder);
      core::LLVMRunPassManager(pass_manager, self.module);
    }
  }

  pub fn new_builder(&self) -> Builder
  {
    Builder(unsafe { core::LLVMCreateBuilderInContext(self.ctx) })
  }
}

impl Drop for JitCompiler {
  fn drop(&mut self) {
    unsafe {
      core::LLVMDisposeModule(self.module);
      core::LLVMContextDispose(self.ctx);
    }
  }
}

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

  #[test]
  fn test_ctor() {
    let jit = JitCompiler::new("target/test-ir/test-module.bc").ok().unwrap();
  }
}<|endoftext|>


<commit_before>use cursive::traits::*;
use cursive::views;

use std::io::Read as _;
use std::io::Write as _;
use std::sync::{Arc, Mutex};

// This example builds a simple TCP server with some parameters and some output.
// It then builds a TUI to control the parameters and display the output.

fn main() {
    let mut siv = cursive::Cursive::default();

    // Build a shared model
    let model = Arc::new(Mutex::new(Model {
        offset: 0,
        logs: Vec::new(),
        cb_sink: siv.cb_sink().clone(),
    }));

    // Start the TCP server in a thread
    start_server(Arc::clone(&model));

    // Build the UI from the model
    siv.add_layer(
        views::Dialog::around(build_ui(Arc::clone(&model)))
            .button("Quit", |s| s.quit()),
    );

    siv.run();
}

struct Model {
    offset: u8,
    logs: Vec<(u8, u8)>,
    cb_sink: cursive::CbSink,
}

fn start_server(model: Arc<Mutex<Model>>) {
    std::thread::spawn(move || {
        if let Err(err) = serve(Arc::clone(&model)) {
            model
                .lock()
                .unwrap()
                .cb_sink
                .send(Box::new(move |s: &mut cursive::Cursive| {
                    s.add_layer(
                        views::Dialog::text(format!("{:?}", err))
                            .title("Error in TCP server")
                            .button("Quit", |s| s.quit()),
                    );
                }))
                .unwrap();
        }
    });
}

fn serve(model: Arc<Mutex<Model>>) -> std::io::Result<()> {
    let listener = std::net::TcpListener::bind("localhost:1234")?;

    for stream in listener.incoming() {
        let stream = stream?;

        for byte in (&stream).bytes() {
            let byte = byte?;
            let mut model = model.lock().unwrap();
            let response = byte.wrapping_add(model.offset);
            model.logs.push((byte, response));
            (&stream).write_all(&[response])?;
            model
                .cb_sink
                .send(Box::new(cursive::Cursive::noop))
                .unwrap();
        }
    }

    Ok(())
}

fn readable_char(byte: u8) -> char {
    if byte.is_ascii_control() {
        '�'
    } else {
        byte as char
    }
}

fn build_log_viewer(model: Arc<Mutex<Model>>) -> impl cursive::view::View {
    views::Canvas::new(model)
        .with_draw(|model, printer| {
            let model = model.lock().unwrap();
            for (i, &(byte, answer)) in model.logs.iter().enumerate() {
                printer.print(
                    (0, i),
                    &format!(
                        "{:3} '{}' -> {:3} '{}'",
                        byte,
                        readable_char(byte),
                        answer,
                        readable_char(answer),
                    ),
                );
            }
        })
        .with_required_size(|model, _req| {
            let model = model.lock().unwrap();
            cursive::Vec2::new(10, model.logs.len())
        })
}

fn build_selector(model: Arc<Mutex<Model>>) -> impl cursive::view::View {
    views::LinearLayout::horizontal()
        .child(
            views::EditView::new()
                .content("0")
                .with_id("edit")
                .min_width(5),
        )
        .child(views::DummyView.fixed_width(1))
        .child(views::Button::new("Update", move |s| {
            if let Some(n) = s
                .call_on_id("edit", |edit: &mut views::EditView| {
                    edit.get_content()
                })
                .and_then(|content| content.parse().ok())
            {
                model.lock().unwrap().offset = n;
            } else {
                s.add_layer(views::Dialog::info(
                    "Could not parse offset as u8",
                ));
            }
        }))
        .child(views::DummyView.fixed_width(1))
        .child(views::Button::new("Test", |s| {
            if let Err(err) = test_server() {
                s.add_layer(
                    views::Dialog::info(format!("{:?}", err))
                        .title("Error running test."),
                );
            }
        }))
}

fn test_server() -> std::io::Result<()> {
    let mut stream = std::net::TcpStream::connect("localhost:1234")?;
    for &byte in &[1, 2, 3, b'a', b'c', b'd'] {
        let mut buf = [0];
        stream.write_all(&[byte])?;
        stream.read_exact(&mut buf)?;
    }
    Ok(())
}

fn build_ui(model: Arc<Mutex<Model>>) -> impl cursive::view::View {
    views::LinearLayout::vertical()
        .child(build_selector(Arc::clone(&model)))
        .child(build_log_viewer(Arc::clone(&model)))
}
<commit_msg>Improve tcp_server example a bit<commit_after>use cursive::traits::*;
use cursive::views;

use std::io::{Read as _, Write as _};
use std::sync::{Arc, Mutex};

// This example builds a simple TCP server with some parameters and some output.
// It then builds a TUI to control the parameters and display the output.

fn main() {
    let mut siv = cursive::Cursive::default();

    // Build a shared model
    let model = Arc::new(Mutex::new(ModelData {
        offset: 10,
        logs: Vec::new(),
        cb_sink: siv.cb_sink().clone(),
    }));

    // Start the TCP server in a thread
    start_server(Arc::clone(&model));

    // Build the UI from the model
    siv.add_layer(
        views::Dialog::around(build_ui(Arc::clone(&model)))
            .button("Quit", |s| s.quit()),
    );

    siv.run();
}

struct ModelData {
    /// The offset will be controlled by the UI and used in the server
    offset: u8,
    /// Logs will be filled by the server and displayed on the UI
    logs: Vec<LogEntry>,
    /// A callback sink is used to control the UI from the server
    /// (eg. force refresh, error popups)
    cb_sink: cursive::CbSink,
}

// Here we use a single mutex, but bigger models might
// prefer individual mutexes for different variables.
type Model = Arc<Mutex<ModelData>>;

#[derive(Clone, Copy)]
struct LogEntry {
    input: u8,
    output: u8,
}

/// Starts serving on a separate thread, and show a popup on error.
fn start_server(model: Model) {
    std::thread::spawn(move || {
        if let Err(err) = serve(Arc::clone(&model)) {
            let model = model.lock().unwrap();
            model
                .cb_sink
                .send(Box::new(move |s: &mut cursive::Cursive| {
                    s.add_layer(
                        views::Dialog::text(format!("{:?}", err))
                            .title("Error in TCP server")
                            .button("Quit", |s| s.quit()),
                    );
                }))
                .unwrap();
        }
    });
}

/// Starts a simple, single-threaded TCP server.
/// Adds a configurable offset to each byte received and sent it back.
fn serve(model: Model) -> std::io::Result<()> {
    // Bind on some local address
    let listener = std::net::TcpListener::bind("localhost:1234")?;

    // Handle each connection sequentially
    for stream in listener.incoming() {
        let stream = stream?;

        // Process each byte according to the current model.
        for byte in (&stream).bytes() {
            let byte = byte?;
            let mut model = model.lock().unwrap();
            let response = byte.wrapping_add(model.offset);
            (&stream).write_all(&[response])?;

            // Save processed jobs
            model.logs.push(LogEntry {
                input: byte,
                output: response,
            });

            // Send a noop to refresh the display
            model
                .cb_sink
                .send(Box::new(cursive::Cursive::noop))
                .unwrap();
        }
    }

    Ok(())
}

/// Build the UI for the given model.
fn build_ui(model: Model) -> impl cursive::view::View {
    // Build the UI in 3 parts, stacked together in a LinearLayout.
    views::LinearLayout::vertical()
        .child(build_selector(Arc::clone(&model)))
        .child(build_tester(Arc::clone(&model)))
        .child(views::DummyView.fixed_height(1))
        .child(build_log_viewer(Arc::clone(&model)))
}

/// Build a view that shows processed jobs from the model.
fn build_log_viewer(model: Model) -> impl cursive::view::View {
    views::Canvas::new(model)
        .with_draw(|model, printer| {
            let model = model.lock().unwrap();
            for (i, &log) in model.logs.iter().enumerate() {
                printer.print(
                    (0, i),
                    &format!(
                        "{:3} '{}'  ->  {:3} '{}'",
                        log.input,
                        readable_char(log.input),
                        log.output,
                        readable_char(log.output),
                    ),
                );
            }
        })
        .with_required_size(|model, _req| {
            let model = model.lock().unwrap();
            cursive::Vec2::new(20, model.logs.len())
        })
        .scrollable()
}

/// Pretty print an ascii u8 if possible.
fn readable_char(byte: u8) -> char {
    if byte.is_ascii_control() {
        '�'
    } else {
        byte as char
    }
}

/// Build a view that can update the model.
fn build_selector(model: Model) -> impl cursive::view::View {
    let offset = model.lock().unwrap().offset;
    views::LinearLayout::horizontal()
        .child(
            views::EditView::new()
                .content(format!("{}", offset))
                .with_id("edit")
                .min_width(5),
        )
        .child(views::DummyView.fixed_width(1))
        .child(views::Button::new("Update", move |s| {
            if let Some(n) = s
                .call_on_id("edit", |edit: &mut views::EditView| {
                    edit.get_content()
                })
                .and_then(|content| content.parse().ok())
            {
                model.lock().unwrap().offset = n;
            } else {
                s.add_layer(views::Dialog::info(
                    "Could not parse offset as u8",
                ));
            }
        }))
}

/// Build a view that can run test connections.
fn build_tester(model: Model) -> impl cursive::view::View {
    views::LinearLayout::horizontal()
        .child(views::TextView::new("Current value:"))
        .child(views::DummyView.fixed_width(1))
        .child(
            views::Canvas::new(model)
                .with_draw(|model, printer| {
                    printer.print(
                        (0, 0),
                        &format!("{}", model.lock().unwrap().offset),
                    )
                })
                .with_required_size(|_, _| cursive::Vec2::new(3, 1)),
        )
        .child(views::DummyView.fixed_width(1))
        .child(views::Button::new("Test", |s| {
            if let Err(err) = test_server() {
                s.add_layer(
                    views::Dialog::info(format!("{:?}", err))
                        .title("Error running test."),
                );
            }
        }))
}

/// Run a test connection.
fn test_server() -> std::io::Result<()> {
    let mut stream = std::net::TcpStream::connect("localhost:1234")?;
    for &byte in b"cursive123" {
        let mut buf = [0];
        stream.write_all(&[byte])?;
        stream.read_exact(&mut buf)?;
    }
    Ok(())
}
<|endoftext|>


<commit_before>//! # OpenTelemetry OTLP Exporter
//!
//! The OpenTelemetry OTLP Exporter supports exporting of trace and metric data in the OTLP format.
mod proto;
mod span;
mod transform;

pub use crate::span::Exporter;
<commit_msg>Export OTLP types (#175)<commit_after>//! # OpenTelemetry OTLP Exporter
//!
//! The OpenTelemetry OTLP Exporter supports exporting of trace and metric data in the OTLP format.
mod proto;
mod span;
mod transform;

pub use crate::span::{Compression, Credentials, Exporter, ExporterConfig, Protocol};
<|endoftext|>


<commit_before>//! Implements a command for uploading proguard mapping files.
use std::fs;
use std::io;
use std::path::PathBuf;

use clap::{App, Arg, ArgMatches};
use console::style;
use failure::{bail, Error, SyncFailure};
use symbolic::common::byteview::ByteView;
use symbolic::proguard::ProguardMappingView;
use uuid::Uuid;

use crate::api::{Api, AssociateDsyms};
use crate::config::Config;
use crate::utils::android::{dump_proguard_uuids_as_properties, AndroidManifest};
use crate::utils::args::{validate_uuid, ArgExt};
use crate::utils::fs::{get_sha1_checksum, TempFile};
use crate::utils::system::QuietExit;
use crate::utils::ui::{copy_with_progress, make_byte_progress_bar};

#[derive(Debug)]
struct MappingRef {
    pub path: PathBuf,
    pub size: u64,
    pub uuid: Uuid,
}

pub fn make_app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
    app.about("Upload ProGuard mapping files to a project.")
        .org_project_args()
        .arg(
            Arg::with_name("paths")
                .value_name("PATH")
                .help("The path to the mapping files.")
                .multiple(true)
                .number_of_values(1)
                .index(1),
        )
        .arg(
            Arg::with_name("version")
                .long("version")
                .value_name("VERSION")
                .requires("app_id")
                .help(
                    "Optionally associate the mapping files with a human \
                     readable version.{n}This helps you understand which \
                     ProGuard files go with which version of your app.",
                ),
        )
        .arg(
            Arg::with_name("version_code")
                .long("version-code")
                .value_name("VERSION_CODE")
                .requires("app_id")
                .requires("version")
                .help(
                    "Optionally associate the mapping files with a version \
                     code.{n}This helps you understand which ProGuard files \
                     go with which version of your app.",
                ),
        )
        .arg(
            Arg::with_name("app_id")
                .long("app-id")
                .value_name("APP_ID")
                .requires("version")
                .help(
                    "Optionally associate the mapping files with an application \
                     ID.{n}If you have multiple apps in one sentry project you can \
                     then easlier tell them apart.",
                ),
        )
        .arg(
            Arg::with_name("platform")
                .long("platform")
                .value_name("PLATFORM")
                .requires("app_id")
                .help(
                    "Optionally defines the platform for the app association. \
                     [defaults to 'android']",
                ),
        )
        .arg(
            Arg::with_name("no_reprocessing")
                .long("no-reprocessing")
                .help("Do not trigger reprocessing after upload."),
        )
        .arg(Arg::with_name("no_upload").long("no-upload").help(
            "Disable the actual upload.{n}This runs all steps for the \
             processing but does not trigger the upload (this also \
             automatically disables reprocessing.  This is useful if you \
             just want to verify the mapping files and write the \
             proguard UUIDs into a properties file.",
        ))
        .arg(
            Arg::with_name("android_manifest")
                .long("android-manifest")
                .value_name("PATH")
                .conflicts_with("app_id")
                .help("Read version and version code from an Android manifest file."),
        )
        .arg(
            Arg::with_name("write_properties")
                .long("write-properties")
                .value_name("PATH")
                .help(
                    "Write the UUIDs for the processed mapping files into \
                     the given properties file.",
                ),
        )
        .arg(
            Arg::with_name("require_one")
                .long("require-one")
                .help("Requires at least one file to upload or the command will error."),
        )
        .arg(
            Arg::with_name("uuid")
                .long("uuid")
                .short("u")
                .value_name("UUID")
                .validator(validate_uuid)
                .help(
                    "Explicitly override the UUID of the mapping file with another one.{n}\
                     This should be used with caution as it means that you can upload \
                     multiple mapping files if you don't take care.  This however can \
                     be useful if you have a build process in which you need to know \
                     the UUID of the proguard file before it was created.  If you upload \
                     a file with a forced UUID you can only upload a single proguard file.",
                ),
        )
}

pub fn execute<'a>(matches: &ArgMatches<'a>) -> Result<(), Error> {
    let api = Api::get_current();

    let paths: Vec<_> = match matches.values_of("paths") {
        Some(paths) => paths.collect(),
        None => {
            return Ok(());
        }
    };
    let mut mappings = vec![];
    let mut all_checksums = vec![];

    let android_manifest = if let Some(path) = matches.value_of("android_manifest") {
        Some(AndroidManifest::from_path(path)?)
    } else {
        None
    };

    let forced_uuid = matches.value_of("uuid").map(|x| x.parse::<Uuid>().unwrap());
    if forced_uuid.is_some() && paths.len() != 1 {
        bail!(
            "When forcing a UUID a single proguard file needs to be \
             provided, got {}",
            paths.len()
        );
    }

    // since the mappings are quite small we don't bother doing a second http
    // request to figure out if any of the checksums are missing.  We just ship
    // them all up.
    for path in &paths {
        match fs::metadata(path) {
            Ok(md) => {
                let byteview = ByteView::from_path(path).map_err(SyncFailure::new)?;
                let mapping = ProguardMappingView::parse(byteview).map_err(SyncFailure::new)?;
                if !mapping.has_line_info() {
                    eprintln!(
                        "warning: proguard mapping '{}' was ignored because it \
                         does not contain any line information.",
                        path
                    );
                } else {
                    let mut f = fs::File::open(path)?;
                    all_checksums.push(get_sha1_checksum(&mut f)?.to_string());
                    mappings.push(MappingRef {
                        path: PathBuf::from(path),
                        size: md.len(),
                        uuid: forced_uuid.unwrap_or_else(|| mapping.uuid()),
                    });
                }
            }
            Err(ref err) if err.kind() == io::ErrorKind::NotFound => {
                eprintln!(
                    "warning: proguard mapping '{}' does not exist. This \
                     might be because the build process did not generate \
                     one (for instance because -dontobfuscate is used)",
                    path
                );
            }
            Err(err) => {
                return Err(Error::from(err)
                    .context(format!("failed to open proguard mapping '{}'", path))
                    .into());
            }
        }
    }

    if mappings.is_empty() && matches.is_present("require_one") {
        println!();
        eprintln!("{}", style("error: found no mapping files to upload").red());
        return Err(QuietExit(1).into());
    }

    println!("{} compressing mappings", style(">").dim());
    let tf = TempFile::create()?;

    // add a scope here so we will flush before uploading
    {
        let mut zip = zip::ZipWriter::new(tf.open());
        for mapping in &mappings {
            let pb = make_byte_progress_bar(mapping.size);
            zip.start_file(
                format!("proguard/{}.txt", mapping.uuid),
                zip::write::FileOptions::default(),
            )?;
            copy_with_progress(&pb, &mut fs::File::open(&mapping.path)?, &mut zip)?;
            pb.finish_and_clear();
        }
    }

    // write UUIDs into the mapping file.
    if let Some(p) = matches.value_of("write_properties") {
        let uuids: Vec<_> = mappings.iter().map(|x| x.uuid).collect();
        dump_proguard_uuids_as_properties(p, &uuids)?;
    }

    if matches.is_present("no_upload") {
        println!("{} skipping upload.", style(">").dim());
        return Ok(());
    }

    println!("{} uploading mappings", style(">").dim());
    let config = Config::get_current();
    let (org, project) = config.get_org_and_project(matches)?;
    let rv = api.upload_dif_archive(&org, &project, tf.path())?;
    println!(
        "{} Uploaded a total of {} new mapping files",
        style(">").dim(),
        style(rv.len()).yellow()
    );
    if !rv.is_empty() {
        println!("Newly uploaded debug symbols:");
        for df in rv {
            println!("  {}", style(&df.id()).dim());
        }
    }

    // update the uuids
    if let Some(android_manifest) = android_manifest {
        api.associate_android_proguard_mappings(&org, &project, &android_manifest, all_checksums)?;

    // if values are given associate
    } else if let Some(app_id) = matches.value_of("app_id") {
        api.associate_dsyms(
            &org,
            &project,
            &AssociateDsyms {
                platform: matches
                    .value_of("platform")
                    .unwrap_or("android")
                    .to_string(),
                checksums: all_checksums,
                name: app_id.to_string(),
                app_id: app_id.to_string(),
                version: matches.value_of("version").unwrap().to_string(),
                build: matches.value_of("version_code").map(|x| x.to_string()),
            },
        )?;
    }

    // If wanted trigger reprocessing
    if !matches.is_present("no_reprocessing") && !matches.is_present("no_upload") {
        if !api.trigger_reprocessing(&org, &project)? {
            println!(
                "{} Server does not support reprocessing. Not triggering.",
                style(">").dim()
            );
        }
    } else {
        println!("{} skipped reprocessing", style(">").dim());
    }

    Ok(())
}
<commit_msg>feat: log proguard mapping file used and hash<commit_after>//! Implements a command for uploading proguard mapping files.
use std::fs;
use std::io;
use std::path::PathBuf;

use clap::{App, Arg, ArgMatches};
use console::style;
use failure::{bail, Error, SyncFailure};
use log::debug;
use symbolic::common::byteview::ByteView;
use symbolic::proguard::ProguardMappingView;
use uuid::Uuid;

use crate::api::{Api, AssociateDsyms};
use crate::config::Config;
use crate::utils::android::{dump_proguard_uuids_as_properties, AndroidManifest};
use crate::utils::args::{validate_uuid, ArgExt};
use crate::utils::fs::{get_sha1_checksum, TempFile};
use crate::utils::system::QuietExit;
use crate::utils::ui::{copy_with_progress, make_byte_progress_bar};

#[derive(Debug)]
struct MappingRef {
    pub path: PathBuf,
    pub size: u64,
    pub uuid: Uuid,
}

pub fn make_app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
    app.about("Upload ProGuard mapping files to a project.")
        .org_project_args()
        .arg(
            Arg::with_name("paths")
                .value_name("PATH")
                .help("The path to the mapping files.")
                .multiple(true)
                .number_of_values(1)
                .index(1),
        )
        .arg(
            Arg::with_name("version")
                .long("version")
                .value_name("VERSION")
                .requires("app_id")
                .help(
                    "Optionally associate the mapping files with a human \
                     readable version.{n}This helps you understand which \
                     ProGuard files go with which version of your app.",
                ),
        )
        .arg(
            Arg::with_name("version_code")
                .long("version-code")
                .value_name("VERSION_CODE")
                .requires("app_id")
                .requires("version")
                .help(
                    "Optionally associate the mapping files with a version \
                     code.{n}This helps you understand which ProGuard files \
                     go with which version of your app.",
                ),
        )
        .arg(
            Arg::with_name("app_id")
                .long("app-id")
                .value_name("APP_ID")
                .requires("version")
                .help(
                    "Optionally associate the mapping files with an application \
                     ID.{n}If you have multiple apps in one sentry project you can \
                     then easlier tell them apart.",
                ),
        )
        .arg(
            Arg::with_name("platform")
                .long("platform")
                .value_name("PLATFORM")
                .requires("app_id")
                .help(
                    "Optionally defines the platform for the app association. \
                     [defaults to 'android']",
                ),
        )
        .arg(
            Arg::with_name("no_reprocessing")
                .long("no-reprocessing")
                .help("Do not trigger reprocessing after upload."),
        )
        .arg(Arg::with_name("no_upload").long("no-upload").help(
            "Disable the actual upload.{n}This runs all steps for the \
             processing but does not trigger the upload (this also \
             automatically disables reprocessing.  This is useful if you \
             just want to verify the mapping files and write the \
             proguard UUIDs into a properties file.",
        ))
        .arg(
            Arg::with_name("android_manifest")
                .long("android-manifest")
                .value_name("PATH")
                .conflicts_with("app_id")
                .help("Read version and version code from an Android manifest file."),
        )
        .arg(
            Arg::with_name("write_properties")
                .long("write-properties")
                .value_name("PATH")
                .help(
                    "Write the UUIDs for the processed mapping files into \
                     the given properties file.",
                ),
        )
        .arg(
            Arg::with_name("require_one")
                .long("require-one")
                .help("Requires at least one file to upload or the command will error."),
        )
        .arg(
            Arg::with_name("uuid")
                .long("uuid")
                .short("u")
                .value_name("UUID")
                .validator(validate_uuid)
                .help(
                    "Explicitly override the UUID of the mapping file with another one.{n}\
                     This should be used with caution as it means that you can upload \
                     multiple mapping files if you don't take care.  This however can \
                     be useful if you have a build process in which you need to know \
                     the UUID of the proguard file before it was created.  If you upload \
                     a file with a forced UUID you can only upload a single proguard file.",
                ),
        )
}

pub fn execute<'a>(matches: &ArgMatches<'a>) -> Result<(), Error> {
    let api = Api::get_current();

    let paths: Vec<_> = match matches.values_of("paths") {
        Some(paths) => paths.collect(),
        None => {
            return Ok(());
        }
    };
    let mut mappings = vec![];
    let mut all_checksums = vec![];

    let android_manifest = if let Some(path) = matches.value_of("android_manifest") {
        Some(AndroidManifest::from_path(path)?)
    } else {
        None
    };

    let forced_uuid = matches.value_of("uuid").map(|x| x.parse::<Uuid>().unwrap());
    if forced_uuid.is_some() && paths.len() != 1 {
        bail!(
            "When forcing a UUID a single proguard file needs to be \
             provided, got {}",
            paths.len()
        );
    }

    // since the mappings are quite small we don't bother doing a second http
    // request to figure out if any of the checksums are missing.  We just ship
    // them all up.
    for path in &paths {
        match fs::metadata(path) {
            Ok(md) => {
                let byteview = ByteView::from_path(path).map_err(SyncFailure::new)?;
                let mapping = ProguardMappingView::parse(byteview).map_err(SyncFailure::new)?;
                if !mapping.has_line_info() {
                    eprintln!(
                        "warning: proguard mapping '{}' was ignored because it \
                         does not contain any line information.",
                        path
                    );
                } else {
                    let mut f = fs::File::open(path)?;
                    let sha = get_sha1_checksum(&mut f)?.to_string();
                    debug!("SHA1 for mapping file '{}': '{}'", path, sha);
                    all_checksums.push(sha);
                    mappings.push(MappingRef {
                        path: PathBuf::from(path),
                        size: md.len(),
                        uuid: forced_uuid.unwrap_or_else(|| mapping.uuid()),
                    });
                }
            }
            Err(ref err) if err.kind() == io::ErrorKind::NotFound => {
                eprintln!(
                    "warning: proguard mapping '{}' does not exist. This \
                     might be because the build process did not generate \
                     one (for instance because -dontobfuscate is used)",
                    path
                );
            }
            Err(err) => {
                return Err(Error::from(err)
                    .context(format!("failed to open proguard mapping '{}'", path))
                    .into());
            }
        }
    }

    if mappings.is_empty() && matches.is_present("require_one") {
        println!();
        eprintln!("{}", style("error: found no mapping files to upload").red());
        return Err(QuietExit(1).into());
    }

    println!("{} compressing mappings", style(">").dim());
    let tf = TempFile::create()?;

    // add a scope here so we will flush before uploading
    {
        let mut zip = zip::ZipWriter::new(tf.open());
        for mapping in &mappings {
            let pb = make_byte_progress_bar(mapping.size);
            zip.start_file(
                format!("proguard/{}.txt", mapping.uuid),
                zip::write::FileOptions::default(),
            )?;
            copy_with_progress(&pb, &mut fs::File::open(&mapping.path)?, &mut zip)?;
            pb.finish_and_clear();
        }
    }

    // write UUIDs into the mapping file.
    if let Some(p) = matches.value_of("write_properties") {
        let uuids: Vec<_> = mappings.iter().map(|x| x.uuid).collect();
        dump_proguard_uuids_as_properties(p, &uuids)?;
    }

    if matches.is_present("no_upload") {
        println!("{} skipping upload.", style(">").dim());
        return Ok(());
    }

    println!("{} uploading mappings", style(">").dim());
    let config = Config::get_current();
    let (org, project) = config.get_org_and_project(matches)?;
    let rv = api.upload_dif_archive(&org, &project, tf.path())?;
    println!(
        "{} Uploaded a total of {} new mapping files",
        style(">").dim(),
        style(rv.len()).yellow()
    );
    if !rv.is_empty() {
        println!("Newly uploaded debug symbols:");
        for df in rv {
            println!("  {}", style(&df.id()).dim());
        }
    }

    // update the uuids
    if let Some(android_manifest) = android_manifest {
        api.associate_android_proguard_mappings(&org, &project, &android_manifest, all_checksums)?;

    // if values are given associate
    } else if let Some(app_id) = matches.value_of("app_id") {
        api.associate_dsyms(
            &org,
            &project,
            &AssociateDsyms {
                platform: matches
                    .value_of("platform")
                    .unwrap_or("android")
                    .to_string(),
                checksums: all_checksums,
                name: app_id.to_string(),
                app_id: app_id.to_string(),
                version: matches.value_of("version").unwrap().to_string(),
                build: matches.value_of("version_code").map(|x| x.to_string()),
            },
        )?;
    }

    // If wanted trigger reprocessing
    if !matches.is_present("no_reprocessing") && !matches.is_present("no_upload") {
        if !api.trigger_reprocessing(&org, &project)? {
            println!(
                "{} Server does not support reprocessing. Not triggering.",
                style(">").dim()
            );
        }
    } else {
        println!("{} skipped reprocessing", style(">").dim());
    }

    Ok(())
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (eliza@elizas.website)
//
//  Copyright (c) 2015-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Parsing and loading Executable and Linkable Format (ELF) 32- and 64-bit
//! binaries.
//!
//! For more information on the ELF format, refer to:
//!
//!  + [Wikipedia](https://en.wikipedia.org/wiki/Executable_and_Linkable_Format)
//!  + The [OS Dev Wiki](http://wiki.osdev.org/ELF)
//!  + The [ELF Format Specification](elfspec)
//!
//! [elfspec]: http://www.skyfree.org/linux/references/ELF_Format.pdf
#![feature(core_intrinsics)]
#![feature(try_from)]
#![no_std]

#[macro_use] extern crate bitflags;
#[macro_use] extern crate macro_attr;

extern crate memory;

use core::{ intrinsics, ops, mem, slice, convert };
use core::convert::TryFrom;

use memory::{ FrameRange, PhysicalPage };

macro_rules! impl_getters {
    ($(#[$attr:meta])* pub fn $name:ident(&self) -> $ty:ty; $($rest:tt)*) => {
        $(#[$attr])* #[inline] pub fn $name(&self) -> $ty { self.$name as $ty }
        impl_getters!{ $( $rest )* }
    };
    ($(#[$attr:meta])* fn $name:ident(&self) -> $ty:ty; $($rest:tt)*) => {
        $(#[$attr])* #[inline] fn $name(&self) -> $ty { self.$name as $ty }
        impl_getters!{ $( $rest )* }
    };
    ( $(#[$attr:meta])* pub fn $name: ident (&self)-> $ty:ty; ) => {
        $(#[$attr])* #[inline] pub fn $name(&self) -> $ty { self.$name as $ty }
    };
    ( $(#[$attr:meta])* fn $name: ident (&self)-> $ty:ty; ) => {
        $(#[$attr])* #[inline] fn $name(&self) -> $ty { self.$name as $ty }
    };
    () => {};
}

pub mod section;
pub mod file;
pub mod program;

/// An ELF section header.
pub type Section<W> = section::Header<Word = W>;
pub type ProgramHeader<W> = program::Header<Word = W>;
/// An ELF header file.
pub type FileHeader<W> = file::HeaderRepr<W>;

/// TODO: should ELF have its own error type?
pub type ElfResult<T> = Result<T, &'static str>;

pub trait ElfWord: Sized + Copy + Clone
                         + ops::Add<Self> + ops::Sub<Self>
                         + ops::Mul<Self> + ops::Div<Self>
                         + ops::Shl<Self> + ops::Shr<Self> { }
impl ElfWord for u64 { }
impl ElfWord for u32 { }

#[cfg(target_pointer_width = "32")]
type DefaultWord = u32;
#[cfg(target_pointer_width = "64")]
type DefaultWord = u64;

/// Hack to make the type-system let me do what I want
trait ValidatesWord<Word: ElfWord> {
    fn check(&self) -> ElfResult<()>;
}

/// A handle on a parsed ELF binary
///  TODO: do we want this to own a HashMap of section names to section headers,
///        to speed up section lookup?
//          - eliza, 03/08/2017
#[derive(Debug)]
pub struct Image< 'bytes             // lifetime of the byte slice
                , Word = DefaultWord // default to machine's pointer size
                , ProgHeader = ProgramHeader<Word> // same word type
                , SectHeader = Section<Word>
                , Header = FileHeader<Word> // must have same word type
                > // jesus christ
where Word: ElfWord + 'bytes
    , ProgHeader: program::Header<Word = Word> + Sized + 'bytes
    , SectHeader: section::Header<Word = Word> + Sized + 'bytes
    , Header: file::Header<Word = Word> + 'bytes
    {
    /// the binary's [file header](file/trait.Header.html)
    pub header: &'bytes Header
  , /// references to each [section header](section/struct.Header.html)
    pub sections: &'bytes [SectHeader]
  , /// references to each [program header](program/trait.Header.html)
    pub program_headers: &'bytes [ProgHeader]
  , /// the raw binary contents of the ELF binary.
    /// note that this includes the _entire_ binary contents of the file,
    /// so the file header and each section header is included in this slice.
    binary: &'bytes [u8]
}

impl <'a, Word, ProgHeader, SectHeader, Header>
Image<'a, Word, ProgHeader, SectHeader, Header>
where Word: ElfWord + 'a
    , ProgHeader: program::Header<Word = Word> + Sized + 'a
    , SectHeader: section::Header<Word = Word> + Sized + 'a
    , Header: file::Header<Word = Word> + 'a
    {
    /// Returns the section header [string table].
    ///
    /// [string table]: section/struct.StrTable.html
    pub fn sh_str_table(&'a self) -> section::StrTable<'a> {
        // TODO: do we want to validate that the string table index is
        //       reasonable (e.g. it's not longer than the binary)?
        //          - eliza, 03/08/2017
        // TODO: do we want to cache a ref to the string table?
        //          - eliza, 03/08/2017
        section::StrTable::from(&self.binary[self.header.sh_str_idx()..])
    }

}

impl<'a, Word, PH, SH, FH> TryFrom<&'a [u8]> for Image<'a, Word, PH, SH, FH>
where Word: ElfWord + 'a
    , PH: program::Header<Word = Word> + 'a
    , SH: section::Header<Word = Word> + 'a
    , FH: file::Header<Word = Word> + 'a
    , &'a FH: convert::TryFrom<&'a [u8], Error = &'static str>
    {

    type Error = &'static str;

    fn try_from(bytes: &'a [u8]) -> ElfResult<Self> {
        let header: &'a FH = <&'a FH>::try_from(bytes)?;

        let sections = unsafe { extract_from_slice::<SH>(
            &bytes[header.sh_range()]
          , 0
          , header.sh_count()
        )? };
        let prog_headers = unsafe { extract_from_slice::<PH>(
            &bytes[header.ph_range()]
          , 0
          , header.ph_count()
        )? };
        Ok(Image { header: header
              , sections: sections
              , program_headers: prog_headers
              , binary: bytes
        })
    }
}

/// Extract `n` instances of type `T` from a byte slice.
///
/// This is essentially just a _slightly_ safer wrapper around
/// [`slice::from_raw_parts`]. Unlike `from_raw_parts`, this function takes
/// a valid byte slice, rather than a pointer. Therefore, some of the safety
/// issues with `from_raw_parts` are avoided:
///
/// + the lifetime (`'slice`) of the returned slice should be the same as the
///   lifetime of the input slice (`data`), rather than inferred arbitrarily.
/// + this function will panic rather than reading past the end of the slice.
///
/// # Arguments
///
/// + `data`: the byte slice to extract a slice of `&[T]`s from
/// + `offset`: a start offset into `data`
/// + `n`: the number of instances of `T` which should be contained
///        in `data[offset..]`
///
/// # Safety
///
/// While this function is safer than [`slice::from_raw_parts`],
/// it is still unsafe for the following reasons:
///
/// + The contents of `data` may not be able to be interpreted as instances of
///   type `T`.
///
/// # Caveats
///
/// + If `n` == 0, this will give you an `&[]`. Just a warning.
//    thanks to Max for making  me figure this out.
/// + `offset` must be aligned on a `T`-sized boundary.
///
/// # Panics
///
/// + If the index `offset` is longer than `T`
///
/// TODO: rewrite this as a `TryFrom` implementation (see issue #85)
//          - eliza, 03/09/2017
///       wait, possibly we should NOT do that. actually we should
///       almost certainly not do that. since this function is unsafe,
///       but `TryFrom` is not, and because this would be WAY generic.
//          - eliza, 03/09/2017
/// TODO: is this general enough to move into util?
//          - eliza, 03/09/2017
/// TODO: refactor this to take a `RangeArgument`?
//          - eliza, 03/13/2017
///       or, we could just remove the offset and expect the caller to
///       offset the slice?
//          - eliza, 03/14/2017
///
/// [`slice::from_raw_parts`]: https://doc.rust-lang.org/stable/std/slice/fn.from_raw_parts.html
unsafe fn extract_from_slice<'slice, T: Sized>( data: &'slice [u8]
                                              , offset: usize
                                              , n: usize)
                                              -> ElfResult<&'slice [T]> {
    if offset % mem::align_of::<T>() != 0 {
        // TODO: these error messages don't contain as much information as they
        //       used to, since the return type is `&'static str` that can't be
        //       dynamically formatted as the panic was. refactor this?
        //       (e.g. should ELF get its own error type?)
        //        - eliza, 03/15/2017
        Err("extract_from_slice: Offset not aligned on type T sized boundary!")
        // assert!(
        //        , "Offset {} not aligned on a {}-sized boundary (must be \
        //           divisible by {})."
        //        , offset, type_name::<T>(), mem::align_of::<T>()
        //        );
    } else if data.len() - offset < mem::size_of::<T>() * n {
        Err("extract_from_slice: Slice too short to contain n instances of T!")
        // assert!(
        //        , "Slice too short to contain {} objects of type {}"
        //        , n, type_name::<T>()
        //        );
    } else {
        Ok(slice::from_raw_parts(data[offset..].as_ptr() as *const T, n))
    }
}

impl<'a, W: ElfWord> convert::Into<FrameRange> for &'a Section<W> {
    #[inline]
    fn into(self) -> FrameRange {
        use memory::PAddr;
        let start = PhysicalPage::from(PAddr::from(self.address() as u64));
        let end = PhysicalPage::from(PAddr::from(self.end_address() as u64));
        start .. end
    }
}
<commit_msg>fix(elf): make converting elf Sections into FrameRanges cross platform<commit_after>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (eliza@elizas.website)
//
//  Copyright (c) 2015-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Parsing and loading Executable and Linkable Format (ELF) 32- and 64-bit
//! binaries.
//!
//! For more information on the ELF format, refer to:
//!
//!  + [Wikipedia](https://en.wikipedia.org/wiki/Executable_and_Linkable_Format)
//!  + The [OS Dev Wiki](http://wiki.osdev.org/ELF)
//!  + The [ELF Format Specification](elfspec)
//!
//! [elfspec]: http://www.skyfree.org/linux/references/ELF_Format.pdf
#![feature(core_intrinsics)]
#![feature(try_from)]
#![no_std]

#[macro_use] extern crate bitflags;
#[macro_use] extern crate macro_attr;

extern crate memory;

use core::{ intrinsics, ops, mem, slice, convert };
use core::convert::TryFrom;

use memory::{ FrameRange, PhysicalPage };

macro_rules! impl_getters {
    ($(#[$attr:meta])* pub fn $name:ident(&self) -> $ty:ty; $($rest:tt)*) => {
        $(#[$attr])* #[inline] pub fn $name(&self) -> $ty { self.$name as $ty }
        impl_getters!{ $( $rest )* }
    };
    ($(#[$attr:meta])* fn $name:ident(&self) -> $ty:ty; $($rest:tt)*) => {
        $(#[$attr])* #[inline] fn $name(&self) -> $ty { self.$name as $ty }
        impl_getters!{ $( $rest )* }
    };
    ( $(#[$attr:meta])* pub fn $name: ident (&self)-> $ty:ty; ) => {
        $(#[$attr])* #[inline] pub fn $name(&self) -> $ty { self.$name as $ty }
    };
    ( $(#[$attr:meta])* fn $name: ident (&self)-> $ty:ty; ) => {
        $(#[$attr])* #[inline] fn $name(&self) -> $ty { self.$name as $ty }
    };
    () => {};
}

pub mod section;
pub mod file;
pub mod program;

/// An ELF section header.
pub type Section<W> = section::Header<Word = W>;
pub type ProgramHeader<W> = program::Header<Word = W>;
/// An ELF header file.
pub type FileHeader<W> = file::HeaderRepr<W>;

/// TODO: should ELF have its own error type?
pub type ElfResult<T> = Result<T, &'static str>;

pub trait ElfWord: Sized + Copy + Clone
                         + ops::Add<Self> + ops::Sub<Self>
                         + ops::Mul<Self> + ops::Div<Self>
                         + ops::Shl<Self> + ops::Shr<Self> { }
impl ElfWord for u64 { }
impl ElfWord for u32 { }

#[cfg(target_pointer_width = "32")]
type DefaultWord = u32;
#[cfg(target_pointer_width = "64")]
type DefaultWord = u64;

/// Hack to make the type-system let me do what I want
trait ValidatesWord<Word: ElfWord> {
    fn check(&self) -> ElfResult<()>;
}

/// A handle on a parsed ELF binary
///  TODO: do we want this to own a HashMap of section names to section headers,
///        to speed up section lookup?
//          - eliza, 03/08/2017
#[derive(Debug)]
pub struct Image< 'bytes             // lifetime of the byte slice
                , Word = DefaultWord // default to machine's pointer size
                , ProgHeader = ProgramHeader<Word> // same word type
                , SectHeader = Section<Word>
                , Header = FileHeader<Word> // must have same word type
                > // jesus christ
where Word: ElfWord + 'bytes
    , ProgHeader: program::Header<Word = Word> + Sized + 'bytes
    , SectHeader: section::Header<Word = Word> + Sized + 'bytes
    , Header: file::Header<Word = Word> + 'bytes
    {
    /// the binary's [file header](file/trait.Header.html)
    pub header: &'bytes Header
  , /// references to each [section header](section/struct.Header.html)
    pub sections: &'bytes [SectHeader]
  , /// references to each [program header](program/trait.Header.html)
    pub program_headers: &'bytes [ProgHeader]
  , /// the raw binary contents of the ELF binary.
    /// note that this includes the _entire_ binary contents of the file,
    /// so the file header and each section header is included in this slice.
    binary: &'bytes [u8]
}

impl <'a, Word, ProgHeader, SectHeader, Header>
Image<'a, Word, ProgHeader, SectHeader, Header>
where Word: ElfWord + 'a
    , ProgHeader: program::Header<Word = Word> + Sized + 'a
    , SectHeader: section::Header<Word = Word> + Sized + 'a
    , Header: file::Header<Word = Word> + 'a
    {
    /// Returns the section header [string table].
    ///
    /// [string table]: section/struct.StrTable.html
    pub fn sh_str_table(&'a self) -> section::StrTable<'a> {
        // TODO: do we want to validate that the string table index is
        //       reasonable (e.g. it's not longer than the binary)?
        //          - eliza, 03/08/2017
        // TODO: do we want to cache a ref to the string table?
        //          - eliza, 03/08/2017
        section::StrTable::from(&self.binary[self.header.sh_str_idx()..])
    }

}

impl<'a, Word, PH, SH, FH> TryFrom<&'a [u8]> for Image<'a, Word, PH, SH, FH>
where Word: ElfWord + 'a
    , PH: program::Header<Word = Word> + 'a
    , SH: section::Header<Word = Word> + 'a
    , FH: file::Header<Word = Word> + 'a
    , &'a FH: convert::TryFrom<&'a [u8], Error = &'static str>
    {

    type Error = &'static str;

    fn try_from(bytes: &'a [u8]) -> ElfResult<Self> {
        let header: &'a FH = <&'a FH>::try_from(bytes)?;

        let sections = unsafe { extract_from_slice::<SH>(
            &bytes[header.sh_range()]
          , 0
          , header.sh_count()
        )? };
        let prog_headers = unsafe { extract_from_slice::<PH>(
            &bytes[header.ph_range()]
          , 0
          , header.ph_count()
        )? };
        Ok(Image { header: header
              , sections: sections
              , program_headers: prog_headers
              , binary: bytes
        })
    }
}

/// Extract `n` instances of type `T` from a byte slice.
///
/// This is essentially just a _slightly_ safer wrapper around
/// [`slice::from_raw_parts`]. Unlike `from_raw_parts`, this function takes
/// a valid byte slice, rather than a pointer. Therefore, some of the safety
/// issues with `from_raw_parts` are avoided:
///
/// + the lifetime (`'slice`) of the returned slice should be the same as the
///   lifetime of the input slice (`data`), rather than inferred arbitrarily.
/// + this function will panic rather than reading past the end of the slice.
///
/// # Arguments
///
/// + `data`: the byte slice to extract a slice of `&[T]`s from
/// + `offset`: a start offset into `data`
/// + `n`: the number of instances of `T` which should be contained
///        in `data[offset..]`
///
/// # Safety
///
/// While this function is safer than [`slice::from_raw_parts`],
/// it is still unsafe for the following reasons:
///
/// + The contents of `data` may not be able to be interpreted as instances of
///   type `T`.
///
/// # Caveats
///
/// + If `n` == 0, this will give you an `&[]`. Just a warning.
//    thanks to Max for making  me figure this out.
/// + `offset` must be aligned on a `T`-sized boundary.
///
/// # Panics
///
/// + If the index `offset` is longer than `T`
///
/// TODO: rewrite this as a `TryFrom` implementation (see issue #85)
//          - eliza, 03/09/2017
///       wait, possibly we should NOT do that. actually we should
///       almost certainly not do that. since this function is unsafe,
///       but `TryFrom` is not, and because this would be WAY generic.
//          - eliza, 03/09/2017
/// TODO: is this general enough to move into util?
//          - eliza, 03/09/2017
/// TODO: refactor this to take a `RangeArgument`?
//          - eliza, 03/13/2017
///       or, we could just remove the offset and expect the caller to
///       offset the slice?
//          - eliza, 03/14/2017
///
/// [`slice::from_raw_parts`]: https://doc.rust-lang.org/stable/std/slice/fn.from_raw_parts.html
unsafe fn extract_from_slice<'slice, T: Sized>( data: &'slice [u8]
                                              , offset: usize
                                              , n: usize)
                                              -> ElfResult<&'slice [T]> {
    if offset % mem::align_of::<T>() != 0 {
        // TODO: these error messages don't contain as much information as they
        //       used to, since the return type is `&'static str` that can't be
        //       dynamically formatted as the panic was. refactor this?
        //       (e.g. should ELF get its own error type?)
        //        - eliza, 03/15/2017
        Err("extract_from_slice: Offset not aligned on type T sized boundary!")
        // assert!(
        //        , "Offset {} not aligned on a {}-sized boundary (must be \
        //           divisible by {})."
        //        , offset, type_name::<T>(), mem::align_of::<T>()
        //        );
    } else if data.len() - offset < mem::size_of::<T>() * n {
        Err("extract_from_slice: Slice too short to contain n instances of T!")
        // assert!(
        //        , "Slice too short to contain {} objects of type {}"
        //        , n, type_name::<T>()
        //        );
    } else {
        Ok(slice::from_raw_parts(data[offset..].as_ptr() as *const T, n))
    }
}

impl<'a, W: ElfWord> convert::Into<FrameRange> for &'a Section<W> {
    #[inline]
    fn into(self) -> FrameRange {
        use memory::{Addr, PAddr};
        // TODO: refactor this disgusting type cast monstrosity hell
        let start = PhysicalPage::from(self.address() as <PAddr as Addr>::Repr);
        let end = PhysicalPage::from(self.end_address() as <PAddr as Addr>::Repr);
        start .. end
    }
}
<|endoftext|>


<commit_before>//! Stream-based consumer implementation.

use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
use std::time::Duration;

use futures::{ready, Stream};
use slab::Slab;

use rdkafka_sys as rdsys;
use rdkafka_sys::types::*;

use crate::client::{ClientContext, NativeClient};
use crate::config::{ClientConfig, FromClientConfig, FromClientConfigAndContext, RDKafkaLogLevel};
use crate::consumer::base_consumer::BaseConsumer;
use crate::consumer::{Consumer, ConsumerContext, DefaultConsumerContext, Rebalance};
use crate::error::{KafkaError, KafkaResult};
use crate::message::BorrowedMessage;
use crate::statistics::Statistics;
use crate::topic_partition_list::TopicPartitionList;
#[cfg(feature = "tokio")]
use crate::util::TokioRuntime;
use crate::util::{AsyncRuntime, NativePtr, Timeout};

/// The [`ConsumerContext`] used by the [`StreamConsumer`]. This context will
/// automatically wake up the message stream when new data is available.
///
/// This type is not intended to be used directly. It will be automatically
/// created by the `StreamConsumer` when necessary.
pub struct StreamConsumerContext<C: ConsumerContext + 'static> {
    inner: C,
    wakers: Arc<Mutex<Slab<Option<Waker>>>>,
}

impl<C: ConsumerContext + 'static> StreamConsumerContext<C> {
    fn new(inner: C) -> StreamConsumerContext<C> {
        StreamConsumerContext {
            inner,
            wakers: Arc::new(Mutex::new(Slab::new())),
        }
    }
}

impl<C: ConsumerContext + 'static> ClientContext for StreamConsumerContext<C> {
    fn log(&self, level: RDKafkaLogLevel, fac: &str, log_message: &str) {
        self.inner.log(level, fac, log_message)
    }

    fn stats(&self, statistics: Statistics) {
        self.inner.stats(statistics)
    }

    fn error(&self, error: KafkaError, reason: &str) {
        self.inner.error(error, reason)
    }
}

impl<C: ConsumerContext + 'static> ConsumerContext for StreamConsumerContext<C> {
    fn rebalance(
        &self,
        native_client: &NativeClient,
        err: RDKafkaRespErr,
        tpl: &mut TopicPartitionList,
    ) {
        self.inner.rebalance(native_client, err, tpl)
    }

    fn pre_rebalance<'a>(&self, rebalance: &Rebalance<'a>) {
        self.inner.pre_rebalance(rebalance)
    }

    fn post_rebalance<'a>(&self, rebalance: &Rebalance<'a>) {
        self.inner.post_rebalance(rebalance)
    }

    fn commit_callback(&self, result: KafkaResult<()>, offsets: &TopicPartitionList) {
        self.inner.commit_callback(result, offsets)
    }

    fn main_queue_min_poll_interval(&self) -> Timeout {
        self.inner.main_queue_min_poll_interval()
    }

    fn message_queue_nonempty_callback(&self) {
        let mut wakers = self.wakers.lock().unwrap();
        for (_, waker) in wakers.iter_mut() {
            if let Some(waker) = waker.take() {
                waker.wake();
            }
        }
        self.inner.message_queue_nonempty_callback()
    }
}

/// A Kafka consumer implementing [`futures::Stream`].
pub struct MessageStream<
    'a,
    C,
    // Ugly, but this provides backwards compatibility when the `tokio` feature
    // is enabled, as it is by default.
    #[cfg(feature = "tokio")] R = TokioRuntime,
    #[cfg(not(feature = "tokio"))] R,
> where
    C: ConsumerContext + 'static,
    R: AsyncRuntime,
{
    consumer: &'a StreamConsumer<C>,
    interval: Duration,
    delay: Pin<Box<Option<R::Delay>>>,
    slot: usize,
}

impl<'a, C, R> MessageStream<'a, C, R>
where
    C: ConsumerContext + 'static,
    R: AsyncRuntime,
{
    fn new(consumer: &'a StreamConsumer<C>, interval: Duration) -> MessageStream<'a, C, R> {
        let slot = {
            let context = consumer.get_base_consumer().context();
            let mut wakers = context.wakers.lock().expect("lock poisoned");
            wakers.insert(None)
        };
        MessageStream {
            consumer,
            interval,
            delay: Box::pin(None),
            slot,
        }
    }

    fn context(&self) -> &StreamConsumerContext<C> {
        self.consumer.get_base_consumer().context()
    }

    fn client_ptr(&self) -> *mut RDKafka {
        self.consumer.client().native_ptr()
    }

    fn set_waker(&self, waker: Waker) {
        let mut wakers = self.context().wakers.lock().expect("lock poisoned");
        wakers[self.slot].replace(waker);
    }

    fn poll(&self) -> Option<KafkaResult<BorrowedMessage<'a>>> {
        unsafe {
            NativePtr::from_ptr(rdsys::rd_kafka_consumer_poll(self.client_ptr(), 0))
                .map(|p| BorrowedMessage::from_consumer(p, self.consumer))
        }
    }

    // SAFETY: All access to `self.delay` occurs via the following two
    // functions. These functions are careful to never move out of `self.delay`.
    // (They can *drop* the future stored in `self.delay`, but that is
    // permitted.) They never return a non-pinned pointer to the contents of
    // `self.delay`.

    fn ensure_delay(&mut self, delay: R::Delay) -> Pin<&mut R::Delay> {
        unsafe { Pin::new_unchecked(self.delay.as_mut().get_unchecked_mut().get_or_insert(delay)) }
    }

    fn clear_delay(&mut self) {
        unsafe { *self.delay.as_mut().get_unchecked_mut() = None }
    }
}

impl<'a, C, R> Stream for MessageStream<'a, C, R>
where
    C: ConsumerContext + 'a,
    R: AsyncRuntime,
{
    type Item = KafkaResult<BorrowedMessage<'a>>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        // Unconditionally store the waker so that we are woken up if the queue
        // flips from non-empty to empty. We have to store the waker on every
        // call to poll in case this future migrates between tasks. We also need
        // to store the waker *before* calling poll to avoid a race where `poll`
        // returns None to indicate that the queue is empty, but the queue
        // becomes non-empty before we've installed the waker.
        self.set_waker(cx.waker().clone());

        match self.poll() {
            None => loop {
                let delay = R::delay_for(self.interval);
                ready!(self.ensure_delay(delay).poll(cx));
                self.clear_delay();
            },
            Some(message) => {
                self.clear_delay();
                Poll::Ready(Some(message))
            }
        }
    }
}

impl<'a, C, R> Drop for MessageStream<'a, C, R>
where
    C: ConsumerContext + 'static,
    R: AsyncRuntime,
{
    fn drop(&mut self) {
        let mut wakers = self.context().wakers.lock().expect("lock poisoned");
        wakers.remove(self.slot);
    }
}

/// A Kafka consumer providing a [`futures::Stream`] interface.
///
/// This consumer doesn't need to be polled explicitly since `await`ing the
/// stream returned by [`StreamConsumer::start`] will implicitly poll the
/// consumer.
#[must_use = "Consumer polling thread will stop immediately if unused"]
pub struct StreamConsumer<C: ConsumerContext + 'static = DefaultConsumerContext> {
    consumer: Arc<BaseConsumer<StreamConsumerContext<C>>>,
}

impl<C: ConsumerContext> Consumer<StreamConsumerContext<C>> for StreamConsumer<C> {
    fn get_base_consumer(&self) -> &BaseConsumer<StreamConsumerContext<C>> {
        Arc::as_ref(&self.consumer)
    }
}

impl FromClientConfig for StreamConsumer {
    fn from_config(config: &ClientConfig) -> KafkaResult<StreamConsumer> {
        StreamConsumer::from_config_and_context(config, DefaultConsumerContext)
    }
}

/// Creates a new `StreamConsumer` starting from a [`ClientConfig`].
impl<C: ConsumerContext> FromClientConfigAndContext<C> for StreamConsumer<C> {
    fn from_config_and_context(
        config: &ClientConfig,
        context: C,
    ) -> KafkaResult<StreamConsumer<C>> {
        let context = StreamConsumerContext::new(context);
        let stream_consumer = StreamConsumer {
            consumer: Arc::new(BaseConsumer::from_config_and_context(config, context)?),
        };
        unsafe {
            rdsys::rd_kafka_poll_set_consumer(stream_consumer.consumer.client().native_ptr())
        };
        Ok(stream_consumer)
    }
}

impl<C: ConsumerContext> StreamConsumer<C> {
    /// Starts the stream consumer with default configuration (100ms polling
    /// interval and no `NoMessageReceived` notifications).
    ///
    /// **Note:** this method must be called from within the context of a Tokio
    /// runtime.
    #[cfg(feature = "tokio")]
    #[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
    pub fn start(&self) -> MessageStream<'_, C, TokioRuntime> {
        self.start_with(Duration::from_millis(100))
    }

    /// Starts the stream consumer with the specified poll interval.
    #[cfg(feature = "tokio")]
    #[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
    pub fn start_with(&self, poll_interval: Duration) -> MessageStream<'_, C, TokioRuntime> {
        // TODO: verify called once
        self.start_with_runtime(poll_interval)
    }

    /// Like [`StreamConsumer::start_with`], but with a customizable
    /// asynchronous runtime.
    ///
    /// See the [`AsyncRuntime`] trait for the details on the interface the
    /// runtime must satisfy.
    pub fn start_with_runtime<R>(&self, poll_interval: Duration) -> MessageStream<'_, C, R>
    where
        R: AsyncRuntime,
    {
        MessageStream::new(self, poll_interval)
    }
}
<commit_msg>Remove unnecessary Arc in StreamConsumer<commit_after>//! Stream-based consumer implementation.

use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
use std::time::Duration;

use futures::{ready, Stream};
use slab::Slab;

use rdkafka_sys as rdsys;
use rdkafka_sys::types::*;

use crate::client::{ClientContext, NativeClient};
use crate::config::{ClientConfig, FromClientConfig, FromClientConfigAndContext, RDKafkaLogLevel};
use crate::consumer::base_consumer::BaseConsumer;
use crate::consumer::{Consumer, ConsumerContext, DefaultConsumerContext, Rebalance};
use crate::error::{KafkaError, KafkaResult};
use crate::message::BorrowedMessage;
use crate::statistics::Statistics;
use crate::topic_partition_list::TopicPartitionList;
#[cfg(feature = "tokio")]
use crate::util::TokioRuntime;
use crate::util::{AsyncRuntime, NativePtr, Timeout};

/// The [`ConsumerContext`] used by the [`StreamConsumer`]. This context will
/// automatically wake up the message stream when new data is available.
///
/// This type is not intended to be used directly. It will be automatically
/// created by the `StreamConsumer` when necessary.
pub struct StreamConsumerContext<C: ConsumerContext + 'static> {
    inner: C,
    wakers: Arc<Mutex<Slab<Option<Waker>>>>,
}

impl<C: ConsumerContext + 'static> StreamConsumerContext<C> {
    fn new(inner: C) -> StreamConsumerContext<C> {
        StreamConsumerContext {
            inner,
            wakers: Arc::new(Mutex::new(Slab::new())),
        }
    }
}

impl<C: ConsumerContext + 'static> ClientContext for StreamConsumerContext<C> {
    fn log(&self, level: RDKafkaLogLevel, fac: &str, log_message: &str) {
        self.inner.log(level, fac, log_message)
    }

    fn stats(&self, statistics: Statistics) {
        self.inner.stats(statistics)
    }

    fn error(&self, error: KafkaError, reason: &str) {
        self.inner.error(error, reason)
    }
}

impl<C: ConsumerContext + 'static> ConsumerContext for StreamConsumerContext<C> {
    fn rebalance(
        &self,
        native_client: &NativeClient,
        err: RDKafkaRespErr,
        tpl: &mut TopicPartitionList,
    ) {
        self.inner.rebalance(native_client, err, tpl)
    }

    fn pre_rebalance<'a>(&self, rebalance: &Rebalance<'a>) {
        self.inner.pre_rebalance(rebalance)
    }

    fn post_rebalance<'a>(&self, rebalance: &Rebalance<'a>) {
        self.inner.post_rebalance(rebalance)
    }

    fn commit_callback(&self, result: KafkaResult<()>, offsets: &TopicPartitionList) {
        self.inner.commit_callback(result, offsets)
    }

    fn main_queue_min_poll_interval(&self) -> Timeout {
        self.inner.main_queue_min_poll_interval()
    }

    fn message_queue_nonempty_callback(&self) {
        let mut wakers = self.wakers.lock().unwrap();
        for (_, waker) in wakers.iter_mut() {
            if let Some(waker) = waker.take() {
                waker.wake();
            }
        }
        self.inner.message_queue_nonempty_callback()
    }
}

/// A Kafka consumer implementing [`futures::Stream`].
pub struct MessageStream<
    'a,
    C,
    // Ugly, but this provides backwards compatibility when the `tokio` feature
    // is enabled, as it is by default.
    #[cfg(feature = "tokio")] R = TokioRuntime,
    #[cfg(not(feature = "tokio"))] R,
> where
    C: ConsumerContext + 'static,
    R: AsyncRuntime,
{
    consumer: &'a StreamConsumer<C>,
    interval: Duration,
    delay: Pin<Box<Option<R::Delay>>>,
    slot: usize,
}

impl<'a, C, R> MessageStream<'a, C, R>
where
    C: ConsumerContext + 'static,
    R: AsyncRuntime,
{
    fn new(consumer: &'a StreamConsumer<C>, interval: Duration) -> MessageStream<'a, C, R> {
        let slot = {
            let context = consumer.get_base_consumer().context();
            let mut wakers = context.wakers.lock().expect("lock poisoned");
            wakers.insert(None)
        };
        MessageStream {
            consumer,
            interval,
            delay: Box::pin(None),
            slot,
        }
    }

    fn context(&self) -> &StreamConsumerContext<C> {
        self.consumer.get_base_consumer().context()
    }

    fn client_ptr(&self) -> *mut RDKafka {
        self.consumer.client().native_ptr()
    }

    fn set_waker(&self, waker: Waker) {
        let mut wakers = self.context().wakers.lock().expect("lock poisoned");
        wakers[self.slot].replace(waker);
    }

    fn poll(&self) -> Option<KafkaResult<BorrowedMessage<'a>>> {
        unsafe {
            NativePtr::from_ptr(rdsys::rd_kafka_consumer_poll(self.client_ptr(), 0))
                .map(|p| BorrowedMessage::from_consumer(p, self.consumer))
        }
    }

    // SAFETY: All access to `self.delay` occurs via the following two
    // functions. These functions are careful to never move out of `self.delay`.
    // (They can *drop* the future stored in `self.delay`, but that is
    // permitted.) They never return a non-pinned pointer to the contents of
    // `self.delay`.

    fn ensure_delay(&mut self, delay: R::Delay) -> Pin<&mut R::Delay> {
        unsafe { Pin::new_unchecked(self.delay.as_mut().get_unchecked_mut().get_or_insert(delay)) }
    }

    fn clear_delay(&mut self) {
        unsafe { *self.delay.as_mut().get_unchecked_mut() = None }
    }
}

impl<'a, C, R> Stream for MessageStream<'a, C, R>
where
    C: ConsumerContext + 'a,
    R: AsyncRuntime,
{
    type Item = KafkaResult<BorrowedMessage<'a>>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        // Unconditionally store the waker so that we are woken up if the queue
        // flips from non-empty to empty. We have to store the waker on every
        // call to poll in case this future migrates between tasks. We also need
        // to store the waker *before* calling poll to avoid a race where `poll`
        // returns None to indicate that the queue is empty, but the queue
        // becomes non-empty before we've installed the waker.
        self.set_waker(cx.waker().clone());

        match self.poll() {
            None => loop {
                let delay = R::delay_for(self.interval);
                ready!(self.ensure_delay(delay).poll(cx));
                self.clear_delay();
            },
            Some(message) => {
                self.clear_delay();
                Poll::Ready(Some(message))
            }
        }
    }
}

impl<'a, C, R> Drop for MessageStream<'a, C, R>
where
    C: ConsumerContext + 'static,
    R: AsyncRuntime,
{
    fn drop(&mut self) {
        let mut wakers = self.context().wakers.lock().expect("lock poisoned");
        wakers.remove(self.slot);
    }
}

/// A Kafka consumer providing a [`futures::Stream`] interface.
///
/// This consumer doesn't need to be polled explicitly since `await`ing the
/// stream returned by [`StreamConsumer::start`] will implicitly poll the
/// consumer.
#[must_use = "Consumer polling thread will stop immediately if unused"]
pub struct StreamConsumer<C: ConsumerContext + 'static = DefaultConsumerContext> {
    consumer: BaseConsumer<StreamConsumerContext<C>>,
}

impl<C: ConsumerContext> Consumer<StreamConsumerContext<C>> for StreamConsumer<C> {
    fn get_base_consumer(&self) -> &BaseConsumer<StreamConsumerContext<C>> {
        &self.consumer
    }
}

impl FromClientConfig for StreamConsumer {
    fn from_config(config: &ClientConfig) -> KafkaResult<StreamConsumer> {
        StreamConsumer::from_config_and_context(config, DefaultConsumerContext)
    }
}

/// Creates a new `StreamConsumer` starting from a [`ClientConfig`].
impl<C: ConsumerContext> FromClientConfigAndContext<C> for StreamConsumer<C> {
    fn from_config_and_context(
        config: &ClientConfig,
        context: C,
    ) -> KafkaResult<StreamConsumer<C>> {
        let context = StreamConsumerContext::new(context);
        let stream_consumer = StreamConsumer {
            consumer: BaseConsumer::from_config_and_context(config, context)?,
        };
        unsafe {
            rdsys::rd_kafka_poll_set_consumer(stream_consumer.consumer.client().native_ptr())
        };
        Ok(stream_consumer)
    }
}

impl<C: ConsumerContext> StreamConsumer<C> {
    /// Starts the stream consumer with default configuration (100ms polling
    /// interval and no `NoMessageReceived` notifications).
    ///
    /// **Note:** this method must be called from within the context of a Tokio
    /// runtime.
    #[cfg(feature = "tokio")]
    #[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
    pub fn start(&self) -> MessageStream<'_, C, TokioRuntime> {
        self.start_with(Duration::from_millis(100))
    }

    /// Starts the stream consumer with the specified poll interval.
    #[cfg(feature = "tokio")]
    #[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
    pub fn start_with(&self, poll_interval: Duration) -> MessageStream<'_, C, TokioRuntime> {
        // TODO: verify called once
        self.start_with_runtime(poll_interval)
    }

    /// Like [`StreamConsumer::start_with`], but with a customizable
    /// asynchronous runtime.
    ///
    /// See the [`AsyncRuntime`] trait for the details on the interface the
    /// runtime must satisfy.
    pub fn start_with_runtime<R>(&self, poll_interval: Duration) -> MessageStream<'_, C, R>
    where
        R: AsyncRuntime,
    {
        MessageStream::new(self, poll_interval)
    }
}
<|endoftext|>


<commit_before>//! Glyph caching

use { freetype, graphics, Texture, TextureSettings };
use std::collections::HashMap;
use graphics::types::Scalar;

extern crate fnv;
use self::fnv::FnvHasher;
use std::hash::BuildHasherDefault;

use std::path::Path;
use error::Error;

pub use graphics::types::FontSize;

/// The type alias for font characters.
pub type Character<'a> = graphics::character::Character<'a, Texture>;

/// A struct used for caching rendered font.
pub struct GlyphCache<'a> {
    /// The font face.
    pub face: freetype::Face<'a>,
    // Maps from fontsize and character to offset, size and texture.
    data: HashMap<(FontSize, char),
                  ([Scalar; 2], [Scalar; 2], Texture),
                  BuildHasherDefault<FnvHasher>>,
}

impl<'a> GlyphCache<'a> {
    /// Constructor for a GlyphCache.
    pub fn new(font: &Path) -> Result<GlyphCache<'static>, Error> {
        let freetype = match freetype::Library::init() {
            Ok(freetype) => freetype,
            Err(why) => return Err(Error::FreetypeError(why)),
        };
        let face = match freetype.new_face(font, 0) {
            Ok(face) => face,
            Err(why) => return Err(Error::FreetypeError(why)),
        };
        Ok(GlyphCache {
            face: face,
            data: HashMap::with_hasher(BuildHasherDefault::<FnvHasher>::default()),
        })
    }

    /// Creates a GlyphCache for a font stored in memory.
    pub fn from_bytes(font: &'a [u8]) -> Result<GlyphCache<'a>, Error> {
        let freetype = match freetype::Library::init() {
            Ok(freetype) => freetype,
            Err(why) => return Err(Error::FreetypeError(why))
        };
        let face = match freetype.new_memory_face(font, 0) {
            Ok(face) => face,
            Err(why) => return Err(Error::FreetypeError(why))
        };
        Ok(GlyphCache {
            face: face,
            data: HashMap::with_hasher(BuildHasherDefault::<FnvHasher>::default())
        })
    }

    /// Get a `Character` from cache, or load it if not there.
    fn get(&mut self, size: FontSize,  ch: char)
    -> &([Scalar; 2], [Scalar; 2], Texture) {
        // Create a `Character` from a given `FontSize` and `char`.
        fn create_character(face: &freetype::Face, size: FontSize, ch: char)
        -> ([Scalar; 2], [Scalar; 2], Texture) {
            face.set_pixel_sizes(0, size).unwrap();
            face.load_char(ch as usize, freetype::face::DEFAULT).unwrap();
            let glyph = face.glyph().get_glyph().unwrap();
            let bitmap_glyph = glyph.to_bitmap(freetype::render_mode::RenderMode::Normal, None)
                .unwrap();
            let bitmap = bitmap_glyph.bitmap();
            let texture = Texture::from_memory_alpha(bitmap.buffer(),
                                                     bitmap.width() as u32,
                                                     bitmap.rows() as u32,
                                                     &TextureSettings::new()).unwrap();
            (
                [bitmap_glyph.left() as f64, bitmap_glyph.top() as f64],
                [(glyph.advance_x() >> 16) as f64, (glyph.advance_y() >> 16) as f64],
                texture,
            )
        }

        let face = &self.face;// necessary to borrow-check
        self.data.entry((size, ch))
                 .or_insert_with(|| create_character(face, size, ch) )
    }

    /// Load all characters in the `chars` iterator for `size`
    pub fn preload_chars<I>(
        &mut self,
        size: FontSize,
        chars: I
    )
        where
            I: Iterator<Item = char>
    {
        for ch in chars {
            self.get(size, ch);
        }
    }

    /// Load all the printable ASCII characters for `size`. Includes space.
    pub fn preload_printable_ascii(&mut self, size: FontSize) {
        // [0x20, 0x7F) contains all printable ASCII characters ([' ', '~'])
        self.preload_chars(size, (0x20u8 .. 0x7F).map(|ch| ch as char));
    }

    /// Return `ch` for `size` if it's already cached. Don't load.
    /// See the `preload_*` functions.
    pub fn opt_character(&self, size: FontSize, ch: char) -> Option<Character> {
        self.data.get(&(size, ch)).map(|&(offset, size, ref texture)| {
            Character {
                offset: offset,
                size: size,
                texture: texture
            }
        })
    }
}

impl<'b> graphics::character::CharacterCache for GlyphCache<'b> {
    type Texture = Texture;

    fn character<'a>(&'a mut self, size: FontSize, ch: char) -> Character<'a> {
        let &(offset, size, ref texture) = self.get(size, ch);
        return Character {
            offset: offset,
            size: size,
            texture: texture
        }
    }
}
<commit_msg>Rewrite GlyphCache constructors to use functional style<commit_after>//! Glyph caching

use { freetype, graphics, Texture, TextureSettings };
use std::collections::HashMap;
use graphics::types::Scalar;

extern crate fnv;
use self::fnv::FnvHasher;
use std::hash::BuildHasherDefault;

use std::path::Path;
use error::Error;

pub use graphics::types::FontSize;

/// The type alias for font characters.
pub type Character<'a> = graphics::character::Character<'a, Texture>;

/// A struct used for caching rendered font.
pub struct GlyphCache<'a> {
    /// The font face.
    pub face: freetype::Face<'a>,
    // Maps from fontsize and character to offset, size and texture.
    data: HashMap<(FontSize, char),
                  ([Scalar; 2], [Scalar; 2], Texture),
                  BuildHasherDefault<FnvHasher>>,
}

impl<'a> GlyphCache<'a> {
    /// Constructor for a GlyphCache.
    pub fn new(font: &Path) -> Result<GlyphCache<'static>, Error> {
        let fnv = BuildHasherDefault::<FnvHasher>::default();
        freetype::Library::init()
                          .and_then(|freetype| freetype.new_face(font, 0) )
                          .map_err( Error::FreetypeError )
                          .map(|face| GlyphCache {
                                          face: face,
                                          data: HashMap::with_hasher(fnv),
                                      } )
    }

    /// Creates a GlyphCache for a font stored in memory.
    pub fn from_bytes(font: &'a [u8]) -> Result<GlyphCache<'a>, Error> {
        let fnv = BuildHasherDefault::<FnvHasher>::default();
        freetype::Library::init()
                          .and_then(|freetype| freetype.new_memory_face(font, 0) )
                          .map_err( Error::FreetypeError )
                          .map(|face| GlyphCache {
                                          face: face,
                                          data: HashMap::with_hasher(fnv),
                                      } )
    }

    /// Get a `Character` from cache, or load it if not there.
    fn get(&mut self, size: FontSize,  ch: char)
    -> &([Scalar; 2], [Scalar; 2], Texture) {
        // Create a `Character` from a given `FontSize` and `char`.
        fn create_character(face: &freetype::Face, size: FontSize, ch: char)
        -> ([Scalar; 2], [Scalar; 2], Texture) {
            face.set_pixel_sizes(0, size).unwrap();
            face.load_char(ch as usize, freetype::face::DEFAULT).unwrap();
            let glyph = face.glyph().get_glyph().unwrap();
            let bitmap_glyph = glyph.to_bitmap(freetype::render_mode::RenderMode::Normal, None)
                .unwrap();
            let bitmap = bitmap_glyph.bitmap();
            let texture = Texture::from_memory_alpha(bitmap.buffer(),
                                                     bitmap.width() as u32,
                                                     bitmap.rows() as u32,
                                                     &TextureSettings::new()).unwrap();
            (
                [bitmap_glyph.left() as f64, bitmap_glyph.top() as f64],
                [(glyph.advance_x() >> 16) as f64, (glyph.advance_y() >> 16) as f64],
                texture,
            )
        }

        let face = &self.face;// necessary to borrow-check
        self.data.entry((size, ch))
                 .or_insert_with(|| create_character(face, size, ch) )
    }

    /// Load all characters in the `chars` iterator for `size`
    pub fn preload_chars<I>(
        &mut self,
        size: FontSize,
        chars: I
    )
        where
            I: Iterator<Item = char>
    {
        for ch in chars {
            self.get(size, ch);
        }
    }

    /// Load all the printable ASCII characters for `size`. Includes space.
    pub fn preload_printable_ascii(&mut self, size: FontSize) {
        // [0x20, 0x7F) contains all printable ASCII characters ([' ', '~'])
        self.preload_chars(size, (0x20u8 .. 0x7F).map(|ch| ch as char));
    }

    /// Return `ch` for `size` if it's already cached. Don't load.
    /// See the `preload_*` functions.
    pub fn opt_character(&self, size: FontSize, ch: char) -> Option<Character> {
        self.data.get(&(size, ch)).map(|&(offset, size, ref texture)| {
            Character {
                offset: offset,
                size: size,
                texture: texture
            }
        })
    }
}

impl<'b> graphics::character::CharacterCache for GlyphCache<'b> {
    type Texture = Texture;

    fn character<'a>(&'a mut self, size: FontSize, ch: char) -> Character<'a> {
        let &(offset, size, ref texture) = self.get(size, ch);
        return Character {
            offset: offset,
            size: size,
            texture: texture
        }
    }
}
<|endoftext|>