<commit_before>use gl;
use std::{fmt, mem, ptr};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use {Display, DisplayImpl, GlObject};
use context::CommandContext;

struct Shader {
    display: Arc<DisplayImpl>,
    id: gl::types::GLuint,
}

impl Drop for Shader {
    fn drop(&mut self) {
        let id = self.id.clone();
        self.display.context.exec(move |: ctxt| {
            unsafe {
                ctxt.gl.DeleteShader(id);
            }
        });
    }
}

/// A combinaison of shaders linked together.
pub struct Program {
    display: Arc<DisplayImpl>,
    #[allow(dead_code)]
    shaders: Vec<Shader>,
    id: gl::types::GLuint,
    uniforms: Arc<HashMap<String, Uniform>>,
    attributes: Arc<HashMap<String, Attribute>>,
    frag_data_locations: Mutex<HashMap<String, Option<u32>>>,
}

/// Informations about a uniform (except its name).
///
/// Internal struct. Not public.
struct Uniform {
    pub location: gl::types::GLint,
    pub ty: gl::types::GLenum,
    pub size: gl::types::GLint,
}

/// Informations about an attribute of a program (except its name).
///
/// Internal struct. Not public.
#[deriving(Show)]
struct Attribute {
    pub location: gl::types::GLint,
    pub ty: gl::types::GLenum,
    pub size: gl::types::GLint,
}

/// Error that can be triggered when creating a `Program`.
#[deriving(Clone, Show)]
pub enum ProgramCreationError {
    /// Error while compiling one of the shaders.
    CompilationError(String),

    /// Error while linking the program.
    LinkingError(String),

    /// One of the request shader type is not supported by the backend.
    ///
    /// Usually the case of geometry shaders.
    ShaderTypeNotSupported,
}

impl ::std::error::Error for ProgramCreationError {
    fn description(&self) -> &str {
        match self {
            &ProgramCreationError::CompilationError(_) => "Compilation error in one of the \
                                                           shaders",
            &ProgramCreationError::LinkingError(_) => "Error while linking shaders together",
            &ProgramCreationError::ShaderTypeNotSupported => "One of the request shader type is \
                                                              not supported by the backend",
        }
    }

    fn detail(&self) -> Option<String> {
        match self {
            &ProgramCreationError::CompilationError(ref s) => Some(s.clone()),
            &ProgramCreationError::LinkingError(ref s) => Some(s.clone()),
            &ProgramCreationError::ShaderTypeNotSupported => None,
        }
    }

    fn cause(&self) -> Option<&::std::error::Error> {
        None
    }
}

impl Program {
    /// Builds a new program.
    ///
    /// A program is a group of shaders linked together.
    ///
    /// # Parameters
    ///
    /// - `vertex_shader`: Source code of the vertex shader.
    /// - `fragment_shader`: Source code of the fragment shader.
    /// - `geometry_shader`: Source code of the geometry shader.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # let display: glium::Display = unsafe { std::mem::uninitialized() };
    /// # let vertex_source = ""; let fragment_source = ""; let geometry_source = "";
    /// let program = glium::Program::new(&display, vertex_source, fragment_source,
    ///     Some(geometry_source));
    /// ```
    ///
    #[experimental = "The list of shaders and the result error will probably change"]
    pub fn new(display: &Display, vertex_shader: &str, fragment_shader: &str,
               geometry_shader: Option<&str>) -> Result<Program, ProgramCreationError>
    {
        let mut shaders_store = Vec::new();
        shaders_store.push(try!(build_shader(display, gl::VERTEX_SHADER, vertex_shader)));
        match geometry_shader {
            Some(gs) => shaders_store.push(try!(build_shader(display, gl::GEOMETRY_SHADER, gs))),
            None => ()
        }
        shaders_store.push(try!(build_shader(display, gl::FRAGMENT_SHADER, fragment_shader)));

        let mut shaders_ids = Vec::new();
        for sh in shaders_store.iter() {
            shaders_ids.push(sh.id);
        }

        let (tx, rx) = channel();
        display.context.context.exec(move |: ctxt| {
            unsafe {
                let id = ctxt.gl.CreateProgram();
                if id == 0 {
                    panic!("glCreateProgram failed");
                }

                // attaching shaders
                for sh in shaders_ids.iter() {
                    ctxt.gl.AttachShader(id, sh.clone());
                }

                // linking and checking for errors
                ctxt.gl.LinkProgram(id);
                {   let mut link_success: gl::types::GLint = mem::uninitialized();
                    ctxt.gl.GetProgramiv(id, gl::LINK_STATUS, &mut link_success);
                    if link_success == 0 {
                        use ProgramCreationError::LinkingError;

                        match ctxt.gl.GetError() {
                            gl::NO_ERROR => (),
                            gl::INVALID_VALUE => {
                                tx.send(Err(LinkingError(format!("glLinkProgram triggered \
                                                                  GL_INVALID_VALUE"))));
                                return;
                            },
                            gl::INVALID_OPERATION => {
                                tx.send(Err(LinkingError(format!("glLinkProgram triggered \
                                                                  GL_INVALID_OPERATION"))));
                                return;
                            },
                            _ => {
                                tx.send(Err(LinkingError(format!("glLinkProgram triggered an \
                                                                  unknown error"))));
                                return;
                            }
                        };

                        let mut error_log_size: gl::types::GLint = mem::uninitialized();
                        ctxt.gl.GetProgramiv(id, gl::INFO_LOG_LENGTH, &mut error_log_size);

                        let mut error_log: Vec<u8> = Vec::with_capacity(error_log_size as uint);
                        ctxt.gl.GetProgramInfoLog(id, error_log_size, &mut error_log_size,
                            error_log.as_mut_slice().as_mut_ptr() as *mut gl::types::GLchar);
                        error_log.set_len(error_log_size as uint);

                        let msg = String::from_utf8(error_log).unwrap();
                        tx.send(Err(LinkingError(msg)));
                        return;
                    }
                }

                tx.send(Ok(id));
            }
        });

        let id = try!(rx.recv());

        let (tx, rx) = channel();
        display.context.context.exec(move |: mut ctxt| {
            unsafe {
                tx.send((
                    reflect_uniforms(&mut ctxt, id),
                    reflect_attributes(&mut ctxt, id)
                ))
            }
        });

        let (uniforms, attributes) = rx.recv();

        Ok(Program {
            display: display.context.clone(),
            shaders: shaders_store,
            id: id,
            uniforms: Arc::new(uniforms),
            attributes: Arc::new(attributes),
            frag_data_locations: Mutex::new(HashMap::new()),
        })
    }

    /// Returns the *location* of an output fragment, if it exists.
    ///
    /// The *location* is a low-level information that is used internally by glium.
    /// You probably don't need to call this function.
    ///
    /// You can declare output fragments in your shaders by writing:
    ///
    /// ```notrust
    /// out vec4 foo;
    /// ```
    ///
    pub fn get_frag_data_location(&self, name: &str) -> Option<u32> {
        // looking for a cached value
        if let Some(result) = self.frag_data_locations.lock().get(name) {
            return result.clone();
        }

        // querying opengl
        let id = self.id.clone();
        let name_c = name.to_c_str();
        let (tx, rx) = channel();
        self.display.context.exec(move |: ctxt| {
            unsafe {
                let value = ctxt.gl.GetFragDataLocation(id, name_c.as_ptr());
                tx.send(value);
            }
        });

        let location = match rx.recv() {
            -1 => None,
            a => Some(a as u32),
        };

        self.frag_data_locations.lock().insert(name.to_string(), location);
        location
    }
}

impl fmt::Show for Program {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        (format!("Program #{}", self.id)).fmt(formatter)
    }
}

impl GlObject for Program {
    fn get_id(&self) -> gl::types::GLuint {
        self.id
    }
}

// TODO: remove this hack
pub fn get_uniforms_locations(program: &Program) -> Arc<HashMap<String, Uniform>>
{
    program.uniforms.clone()
}

// TODO: remove this hack
pub fn get_attributes(program: &Program) -> Arc<HashMap<String, Attribute>>
{
    program.attributes.clone()
}

impl Drop for Program {
    fn drop(&mut self) {
        // removing VAOs which contain this program
        {
            let mut vaos = self.display.vertex_array_objects.lock();
            let to_delete = vaos.keys().filter(|&&(_, _, p)| p == self.id)
                .map(|k| k.clone()).collect::<Vec<_>>();
            for k in to_delete.into_iter() {
                vaos.remove(&k);
            }
        }

        // sending the destroy command
        let id = self.id.clone();
        self.display.context.exec(move |: ctxt| {
            unsafe {
                if ctxt.state.program == id {
                    ctxt.gl.UseProgram(0);
                    ctxt.state.program = 0;
                }

                ctxt.gl.DeleteProgram(id);
            }
        });
    }
}

/// Builds an individual shader.
fn build_shader<S: ToCStr>(display: &Display, shader_type: gl::types::GLenum, source_code: S)
    -> Result<Shader, ProgramCreationError>
{
    let source_code = source_code.to_c_str();

    let (tx, rx) = channel();
    display.context.context.exec(move |: ctxt| {
        unsafe {
            if shader_type == gl::GEOMETRY_SHADER && ctxt.opengl_es {
                tx.send(Err(ProgramCreationError::ShaderTypeNotSupported));
                return;
            }

            let id = ctxt.gl.CreateShader(shader_type);

            if id == 0 {
                tx.send(Err(ProgramCreationError::ShaderTypeNotSupported));
                return;
            }

            ctxt.gl.ShaderSource(id, 1, [ source_code.as_ptr() ].as_ptr(), ptr::null());
            ctxt.gl.CompileShader(id);

            // checking compilation success
            let compilation_success = {
                let mut compilation_success: gl::types::GLint = mem::uninitialized();
                ctxt.gl.GetShaderiv(id, gl::COMPILE_STATUS, &mut compilation_success);
                compilation_success
            };

            if compilation_success == 0 {
                // compilation error
                let mut error_log_size: gl::types::GLint = mem::uninitialized();
                ctxt.gl.GetShaderiv(id, gl::INFO_LOG_LENGTH, &mut error_log_size);

                let mut error_log: Vec<u8> = Vec::with_capacity(error_log_size as uint);
                ctxt.gl.GetShaderInfoLog(id, error_log_size, &mut error_log_size,
                    error_log.as_mut_slice().as_mut_ptr() as *mut gl::types::GLchar);
                error_log.set_len(error_log_size as uint);

                let msg = String::from_utf8(error_log).unwrap();
                tx.send(Err(ProgramCreationError::CompilationError(msg)));
                return;
            }

            tx.send(Ok(id));
        }
    });

    rx.recv().map(|id| {
        Shader {
            display: display.context.clone(),
            id: id
        }
    })
}

unsafe fn reflect_uniforms(ctxt: &mut CommandContext, program: gl::types::GLuint)
    -> HashMap<String, Uniform>
{
    // reflecting program uniforms
    let mut uniforms = HashMap::new();

    let mut active_uniforms: gl::types::GLint = mem::uninitialized();
    ctxt.gl.GetProgramiv(program, gl::ACTIVE_UNIFORMS, &mut active_uniforms);

    for uniform_id in range(0, active_uniforms) {
        let mut uniform_name_tmp: Vec<u8> = Vec::with_capacity(64);
        let mut uniform_name_tmp_len = 63;

        let mut data_type: gl::types::GLenum = mem::uninitialized();
        let mut data_size: gl::types::GLint = mem::uninitialized();
        ctxt.gl.GetActiveUniform(program, uniform_id as gl::types::GLuint, uniform_name_tmp_len,
            &mut uniform_name_tmp_len, &mut data_size, &mut data_type,
            uniform_name_tmp.as_mut_slice().as_mut_ptr() as *mut gl::types::GLchar);
        uniform_name_tmp.set_len(uniform_name_tmp_len as uint);

        let uniform_name = String::from_utf8(uniform_name_tmp).unwrap();
        let location = ctxt.gl.GetUniformLocation(program, uniform_name.to_c_str().into_inner());

        uniforms.insert(uniform_name, Uniform {
            location: location,
            ty: data_type,
            size: data_size
        });
    }

    uniforms
}

unsafe fn reflect_attributes(ctxt: &mut CommandContext, program: gl::types::GLuint)
    -> HashMap<String, Attribute>
{
    let mut attributes = HashMap::new();

    let mut active_attributes: gl::types::GLint = mem::uninitialized();
    ctxt.gl.GetProgramiv(program, gl::ACTIVE_ATTRIBUTES, &mut active_attributes);

    for attribute_id in range(0, active_attributes) {
        let mut attr_name_tmp: Vec<u8> = Vec::with_capacity(64);
        let mut attr_name_tmp_len = 63;

        let mut data_type: gl::types::GLenum = mem::uninitialized();
        let mut data_size: gl::types::GLint = mem::uninitialized();
        ctxt.gl.GetActiveAttrib(program, attribute_id as gl::types::GLuint, attr_name_tmp_len,
            &mut attr_name_tmp_len, &mut data_size, &mut data_type,
            attr_name_tmp.as_mut_slice().as_mut_ptr() as *mut gl::types::GLchar);
        attr_name_tmp.set_len(attr_name_tmp_len as uint);

        let attr_name = String::from_utf8(attr_name_tmp).unwrap();
        let location = ctxt.gl.GetAttribLocation(program, attr_name.to_c_str().into_inner());

        attributes.insert(attr_name, Attribute {
            location: location,
            ty: data_type,
            size: data_size
        });
    }

    attributes
}
<commit_msg>Use a global mutex to avoid race-conditions in shader compiler<commit_after>use gl;
use std::{fmt, mem, ptr};
use std::collections::HashMap;
use std::sync::{Arc, Mutex, StaticMutex, MUTEX_INIT};
use {Display, DisplayImpl, GlObject};
use context::CommandContext;

/// Some shader compilers have race-condition issues.
/// We lock this mutex in the GL thread every time we compile a shader or link a program.
static COMPILER_GLOBAL_LOCK: StaticMutex = MUTEX_INIT;

struct Shader {
    display: Arc<DisplayImpl>,
    id: gl::types::GLuint,
}

impl Drop for Shader {
    fn drop(&mut self) {
        let id = self.id.clone();
        self.display.context.exec(move |: ctxt| {
            unsafe {
                ctxt.gl.DeleteShader(id);
            }
        });
    }
}

/// A combinaison of shaders linked together.
pub struct Program {
    display: Arc<DisplayImpl>,
    #[allow(dead_code)]
    shaders: Vec<Shader>,
    id: gl::types::GLuint,
    uniforms: Arc<HashMap<String, Uniform>>,
    attributes: Arc<HashMap<String, Attribute>>,
    frag_data_locations: Mutex<HashMap<String, Option<u32>>>,
}

/// Informations about a uniform (except its name).
///
/// Internal struct. Not public.
struct Uniform {
    pub location: gl::types::GLint,
    pub ty: gl::types::GLenum,
    pub size: gl::types::GLint,
}

/// Informations about an attribute of a program (except its name).
///
/// Internal struct. Not public.
#[deriving(Show)]
struct Attribute {
    pub location: gl::types::GLint,
    pub ty: gl::types::GLenum,
    pub size: gl::types::GLint,
}

/// Error that can be triggered when creating a `Program`.
#[deriving(Clone, Show)]
pub enum ProgramCreationError {
    /// Error while compiling one of the shaders.
    CompilationError(String),

    /// Error while linking the program.
    LinkingError(String),

    /// One of the request shader type is not supported by the backend.
    ///
    /// Usually the case of geometry shaders.
    ShaderTypeNotSupported,
}

impl ::std::error::Error for ProgramCreationError {
    fn description(&self) -> &str {
        match self {
            &ProgramCreationError::CompilationError(_) => "Compilation error in one of the \
                                                           shaders",
            &ProgramCreationError::LinkingError(_) => "Error while linking shaders together",
            &ProgramCreationError::ShaderTypeNotSupported => "One of the request shader type is \
                                                              not supported by the backend",
        }
    }

    fn detail(&self) -> Option<String> {
        match self {
            &ProgramCreationError::CompilationError(ref s) => Some(s.clone()),
            &ProgramCreationError::LinkingError(ref s) => Some(s.clone()),
            &ProgramCreationError::ShaderTypeNotSupported => None,
        }
    }

    fn cause(&self) -> Option<&::std::error::Error> {
        None
    }
}

impl Program {
    /// Builds a new program.
    ///
    /// A program is a group of shaders linked together.
    ///
    /// # Parameters
    ///
    /// - `vertex_shader`: Source code of the vertex shader.
    /// - `fragment_shader`: Source code of the fragment shader.
    /// - `geometry_shader`: Source code of the geometry shader.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # let display: glium::Display = unsafe { std::mem::uninitialized() };
    /// # let vertex_source = ""; let fragment_source = ""; let geometry_source = "";
    /// let program = glium::Program::new(&display, vertex_source, fragment_source,
    ///     Some(geometry_source));
    /// ```
    ///
    #[experimental = "The list of shaders and the result error will probably change"]
    pub fn new(display: &Display, vertex_shader: &str, fragment_shader: &str,
               geometry_shader: Option<&str>) -> Result<Program, ProgramCreationError>
    {
        let mut shaders_store = Vec::new();
        shaders_store.push(try!(build_shader(display, gl::VERTEX_SHADER, vertex_shader)));
        match geometry_shader {
            Some(gs) => shaders_store.push(try!(build_shader(display, gl::GEOMETRY_SHADER, gs))),
            None => ()
        }
        shaders_store.push(try!(build_shader(display, gl::FRAGMENT_SHADER, fragment_shader)));

        let mut shaders_ids = Vec::new();
        for sh in shaders_store.iter() {
            shaders_ids.push(sh.id);
        }

        let (tx, rx) = channel();
        display.context.context.exec(move |: ctxt| {
            unsafe {
                let id = ctxt.gl.CreateProgram();
                if id == 0 {
                    panic!("glCreateProgram failed");
                }

                // attaching shaders
                for sh in shaders_ids.iter() {
                    ctxt.gl.AttachShader(id, sh.clone());
                }

                // linking
                {
                    let _lock = COMPILER_GLOBAL_LOCK.lock();
                    ctxt.gl.LinkProgram(id);
                }

                // checking for errors
                {   let mut link_success: gl::types::GLint = mem::uninitialized();
                    ctxt.gl.GetProgramiv(id, gl::LINK_STATUS, &mut link_success);
                    if link_success == 0 {
                        use ProgramCreationError::LinkingError;

                        match ctxt.gl.GetError() {
                            gl::NO_ERROR => (),
                            gl::INVALID_VALUE => {
                                tx.send(Err(LinkingError(format!("glLinkProgram triggered \
                                                                  GL_INVALID_VALUE"))));
                                return;
                            },
                            gl::INVALID_OPERATION => {
                                tx.send(Err(LinkingError(format!("glLinkProgram triggered \
                                                                  GL_INVALID_OPERATION"))));
                                return;
                            },
                            _ => {
                                tx.send(Err(LinkingError(format!("glLinkProgram triggered an \
                                                                  unknown error"))));
                                return;
                            }
                        };

                        let mut error_log_size: gl::types::GLint = mem::uninitialized();
                        ctxt.gl.GetProgramiv(id, gl::INFO_LOG_LENGTH, &mut error_log_size);

                        let mut error_log: Vec<u8> = Vec::with_capacity(error_log_size as uint);
                        ctxt.gl.GetProgramInfoLog(id, error_log_size, &mut error_log_size,
                            error_log.as_mut_slice().as_mut_ptr() as *mut gl::types::GLchar);
                        error_log.set_len(error_log_size as uint);

                        let msg = String::from_utf8(error_log).unwrap();
                        tx.send(Err(LinkingError(msg)));
                        return;
                    }
                }

                tx.send(Ok(id));
            }
        });

        let id = try!(rx.recv());

        let (tx, rx) = channel();
        display.context.context.exec(move |: mut ctxt| {
            unsafe {
                tx.send((
                    reflect_uniforms(&mut ctxt, id),
                    reflect_attributes(&mut ctxt, id)
                ))
            }
        });

        let (uniforms, attributes) = rx.recv();

        Ok(Program {
            display: display.context.clone(),
            shaders: shaders_store,
            id: id,
            uniforms: Arc::new(uniforms),
            attributes: Arc::new(attributes),
            frag_data_locations: Mutex::new(HashMap::new()),
        })
    }

    /// Returns the *location* of an output fragment, if it exists.
    ///
    /// The *location* is a low-level information that is used internally by glium.
    /// You probably don't need to call this function.
    ///
    /// You can declare output fragments in your shaders by writing:
    ///
    /// ```notrust
    /// out vec4 foo;
    /// ```
    ///
    pub fn get_frag_data_location(&self, name: &str) -> Option<u32> {
        // looking for a cached value
        if let Some(result) = self.frag_data_locations.lock().get(name) {
            return result.clone();
        }

        // querying opengl
        let id = self.id.clone();
        let name_c = name.to_c_str();
        let (tx, rx) = channel();
        self.display.context.exec(move |: ctxt| {
            unsafe {
                let value = ctxt.gl.GetFragDataLocation(id, name_c.as_ptr());
                tx.send(value);
            }
        });

        let location = match rx.recv() {
            -1 => None,
            a => Some(a as u32),
        };

        self.frag_data_locations.lock().insert(name.to_string(), location);
        location
    }
}

impl fmt::Show for Program {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        (format!("Program #{}", self.id)).fmt(formatter)
    }
}

impl GlObject for Program {
    fn get_id(&self) -> gl::types::GLuint {
        self.id
    }
}

// TODO: remove this hack
pub fn get_uniforms_locations(program: &Program) -> Arc<HashMap<String, Uniform>>
{
    program.uniforms.clone()
}

// TODO: remove this hack
pub fn get_attributes(program: &Program) -> Arc<HashMap<String, Attribute>>
{
    program.attributes.clone()
}

impl Drop for Program {
    fn drop(&mut self) {
        // removing VAOs which contain this program
        {
            let mut vaos = self.display.vertex_array_objects.lock();
            let to_delete = vaos.keys().filter(|&&(_, _, p)| p == self.id)
                .map(|k| k.clone()).collect::<Vec<_>>();
            for k in to_delete.into_iter() {
                vaos.remove(&k);
            }
        }

        // sending the destroy command
        let id = self.id.clone();
        self.display.context.exec(move |: ctxt| {
            unsafe {
                if ctxt.state.program == id {
                    ctxt.gl.UseProgram(0);
                    ctxt.state.program = 0;
                }

                ctxt.gl.DeleteProgram(id);
            }
        });
    }
}

/// Builds an individual shader.
fn build_shader<S: ToCStr>(display: &Display, shader_type: gl::types::GLenum, source_code: S)
    -> Result<Shader, ProgramCreationError>
{
    let source_code = source_code.to_c_str();

    let (tx, rx) = channel();
    display.context.context.exec(move |: ctxt| {
        unsafe {
            if shader_type == gl::GEOMETRY_SHADER && ctxt.opengl_es {
                tx.send(Err(ProgramCreationError::ShaderTypeNotSupported));
                return;
            }

            let id = ctxt.gl.CreateShader(shader_type);

            if id == 0 {
                tx.send(Err(ProgramCreationError::ShaderTypeNotSupported));
                return;
            }

            ctxt.gl.ShaderSource(id, 1, [ source_code.as_ptr() ].as_ptr(), ptr::null());

            // compiling
            {
                let _lock = COMPILER_GLOBAL_LOCK.lock();
                ctxt.gl.CompileShader(id);
            }

            // checking compilation success
            let compilation_success = {
                let mut compilation_success: gl::types::GLint = mem::uninitialized();
                ctxt.gl.GetShaderiv(id, gl::COMPILE_STATUS, &mut compilation_success);
                compilation_success
            };

            if compilation_success == 0 {
                // compilation error
                let mut error_log_size: gl::types::GLint = mem::uninitialized();
                ctxt.gl.GetShaderiv(id, gl::INFO_LOG_LENGTH, &mut error_log_size);

                let mut error_log: Vec<u8> = Vec::with_capacity(error_log_size as uint);
                ctxt.gl.GetShaderInfoLog(id, error_log_size, &mut error_log_size,
                    error_log.as_mut_slice().as_mut_ptr() as *mut gl::types::GLchar);
                error_log.set_len(error_log_size as uint);

                let msg = String::from_utf8(error_log).unwrap();
                tx.send(Err(ProgramCreationError::CompilationError(msg)));
                return;
            }

            tx.send(Ok(id));
        }
    });

    rx.recv().map(|id| {
        Shader {
            display: display.context.clone(),
            id: id
        }
    })
}

unsafe fn reflect_uniforms(ctxt: &mut CommandContext, program: gl::types::GLuint)
    -> HashMap<String, Uniform>
{
    // reflecting program uniforms
    let mut uniforms = HashMap::new();

    let mut active_uniforms: gl::types::GLint = mem::uninitialized();
    ctxt.gl.GetProgramiv(program, gl::ACTIVE_UNIFORMS, &mut active_uniforms);

    for uniform_id in range(0, active_uniforms) {
        let mut uniform_name_tmp: Vec<u8> = Vec::with_capacity(64);
        let mut uniform_name_tmp_len = 63;

        let mut data_type: gl::types::GLenum = mem::uninitialized();
        let mut data_size: gl::types::GLint = mem::uninitialized();
        ctxt.gl.GetActiveUniform(program, uniform_id as gl::types::GLuint, uniform_name_tmp_len,
            &mut uniform_name_tmp_len, &mut data_size, &mut data_type,
            uniform_name_tmp.as_mut_slice().as_mut_ptr() as *mut gl::types::GLchar);
        uniform_name_tmp.set_len(uniform_name_tmp_len as uint);

        let uniform_name = String::from_utf8(uniform_name_tmp).unwrap();
        let location = ctxt.gl.GetUniformLocation(program, uniform_name.to_c_str().into_inner());

        uniforms.insert(uniform_name, Uniform {
            location: location,
            ty: data_type,
            size: data_size
        });
    }

    uniforms
}

unsafe fn reflect_attributes(ctxt: &mut CommandContext, program: gl::types::GLuint)
    -> HashMap<String, Attribute>
{
    let mut attributes = HashMap::new();

    let mut active_attributes: gl::types::GLint = mem::uninitialized();
    ctxt.gl.GetProgramiv(program, gl::ACTIVE_ATTRIBUTES, &mut active_attributes);

    for attribute_id in range(0, active_attributes) {
        let mut attr_name_tmp: Vec<u8> = Vec::with_capacity(64);
        let mut attr_name_tmp_len = 63;

        let mut data_type: gl::types::GLenum = mem::uninitialized();
        let mut data_size: gl::types::GLint = mem::uninitialized();
        ctxt.gl.GetActiveAttrib(program, attribute_id as gl::types::GLuint, attr_name_tmp_len,
            &mut attr_name_tmp_len, &mut data_size, &mut data_type,
            attr_name_tmp.as_mut_slice().as_mut_ptr() as *mut gl::types::GLchar);
        attr_name_tmp.set_len(attr_name_tmp_len as uint);

        let attr_name = String::from_utf8(attr_name_tmp).unwrap();
        let location = ctxt.gl.GetAttribLocation(program, attr_name.to_c_str().into_inner());

        attributes.insert(attr_name, Attribute {
            location: location,
            ty: data_type,
            size: data_size
        });
    }

    attributes
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use autocxx_parser::file_locations::FileLocationStrategy;
use miette::Diagnostic;
use proc_macro2::TokenStream;
use thiserror::Error;

use crate::{strip_system_headers, CppCodegenOptions, ParseError, RebuildDependencyRecorder};
use std::ffi::OsStr;
use std::ffi::OsString;
use std::fs::File;
use std::io::Write;
use std::marker::PhantomData;
use std::path::{Path, PathBuf};

/// Errors returned during creation of a [`cc::Build`] from an include_cxx
/// macro.
#[derive(Error, Diagnostic, Debug)]
#[cfg_attr(feature = "nightly", doc(cfg(feature = "build")))]
pub enum BuilderError {
    #[error("cxx couldn't handle our generated bindings - could be a bug in autocxx: {0}")]
    InvalidCxx(cxx_gen::Error),
    #[error(transparent)]
    #[diagnostic(transparent)]
    ParseError(ParseError),
    #[error("we couldn't write the generated code to disk at {1}: {0}")]
    FileWriteFail(std::io::Error, PathBuf),
    #[error("no include_cpp! macro was found")]
    NoIncludeCxxMacrosFound,
    #[error("could not create a directory {1}: {0}")]
    UnableToCreateDirectory(std::io::Error, PathBuf),
}

#[cfg_attr(feature = "nightly", doc(cfg(feature = "build")))]
pub type BuilderBuild = cc::Build;

/// For test purposes only, a [`cc::Build`] and lists of Rust and C++
/// files generated.
#[cfg_attr(feature = "nightly", doc(cfg(feature = "build")))]
pub struct BuilderSuccess(pub BuilderBuild, pub Vec<PathBuf>, pub Vec<PathBuf>);

/// Results of a build.
#[cfg_attr(feature = "nightly", doc(cfg(feature = "build")))]
pub type BuilderResult = Result<BuilderSuccess, BuilderError>;

/// The context in which a builder object lives. Callbacks for various
/// purposes.
#[cfg_attr(feature = "nightly", doc(cfg(feature = "build")))]
pub trait BuilderContext {
    /// Perform any initialization specific to the context in which this
    /// builder lives.
    fn setup() {}

    /// Create a dependency recorder, if any.
    fn get_dependency_recorder() -> Option<Box<dyn RebuildDependencyRecorder>>;
}

/// An object to allow building of bindings from a `build.rs` file.
///
/// It would be unusual to use this directly - see the `autocxx_build` or
/// `autocxx_gen` crates.
#[cfg_attr(feature = "nightly", doc(cfg(feature = "build")))]
pub struct Builder<'a, BuilderContext> {
    rs_file: PathBuf,
    autocxx_incs: Vec<OsString>,
    extra_clang_args: Vec<String>,
    dependency_recorder: Option<Box<dyn RebuildDependencyRecorder>>,
    custom_gendir: Option<PathBuf>,
    auto_allowlist: bool,
    cpp_codegen_options: CppCodegenOptions<'a>,
    // This member is to ensure that this type is parameterized
    // by a BuilderContext. The goal is to balance three needs:
    // (1) have most of the functionality over in autocxx_engine,
    // (2) expose this type to users of autocxx_build and to
    //     make it easy for callers simply to call Builder::new,
    // (3) ensure that such a Builder does a few tasks specific to its use
    // in a cargo environment.
    ctx: PhantomData<BuilderContext>,
}

impl<CTX: BuilderContext> Builder<'_, CTX> {
    pub fn new(
        rs_file: impl AsRef<Path>,
        autocxx_incs: impl IntoIterator<Item = impl AsRef<OsStr>>,
    ) -> Self {
        CTX::setup();
        Self {
            rs_file: rs_file.as_ref().to_path_buf(),
            autocxx_incs: autocxx_incs
                .into_iter()
                .map(|s| s.as_ref().to_os_string())
                .collect(),
            extra_clang_args: Vec::new(),
            dependency_recorder: CTX::get_dependency_recorder(),
            custom_gendir: None,
            auto_allowlist: false,
            cpp_codegen_options: CppCodegenOptions::default(),
            ctx: PhantomData,
        }
    }

    /// Specify extra arguments for clang.
    pub fn extra_clang_args(mut self, extra_clang_args: &[&str]) -> Self {
        self.extra_clang_args = extra_clang_args.iter().map(|s| s.to_string()).collect();
        self
    }

    /// Where to generate the code.
    pub fn custom_gendir(mut self, custom_gendir: PathBuf) -> Self {
        self.custom_gendir = Some(custom_gendir);
        self
    }

    /// Update C++ code generation options. See [`CppCodegenOptions`] for details.
    pub fn cpp_codegen_options<F>(mut self, modifier: F) -> Self
    where
        F: FnOnce(&mut CppCodegenOptions),
    {
        modifier(&mut self.cpp_codegen_options);
        self
    }

    /// Automatically discover uses of the C++ `ffi` mod and generate the allowlist
    /// from that.
    /// This is a highly experimental option, not currently recommended.
    /// It doesn't work in the following cases:
    /// * Static function calls on types within the FFI mod.
    /// * Anything inside a macro invocation.
    /// * You're using a different name for your `ffi` mod
    /// * You're using multiple FFI mods
    /// * You've got usages scattered across files beyond that with the
    ///   `include_cpp` invocation
    /// * You're using `use` statements to rename mods or items. If this
    /// proves to be a promising or helpful direction, autocxx would be happy
    /// to accept pull requests to remove some of these limitations.
    pub fn auto_allowlist(mut self, do_it: bool) -> Self {
        self.auto_allowlist = do_it;
        self
    }

    /// Whether to suppress inclusion of system headers (`memory`, `string` etc.)
    /// from generated C++ bindings code. This should not normally be used,
    /// but can occasionally be useful if you're reducing a test case and you
    /// have a preprocessed header file which already contains absolutely everything
    /// that the bindings could ever need.
    pub fn suppress_system_headers(mut self, do_it: bool) -> Self {
        self.cpp_codegen_options.suppress_system_headers = do_it;
        self
    }

    /// An annotation optionally to include on each C++ function.
    /// For example to export the symbol from a library.
    pub fn cxx_impl_annotations(mut self, cxx_impl_annotations: Option<String>) -> Self {
        self.cpp_codegen_options.cxx_impl_annotations = cxx_impl_annotations;
        self
    }

    /// Whether to skip using [`cxx_gen`] to generate the C++ code,
    /// so that some other process can handle that.
    pub fn skip_cxx_gen(mut self, skip_cxx_gen: bool) -> Self {
        self.cpp_codegen_options.skip_cxx_gen = skip_cxx_gen;
        self
    }

    /// Build autocxx C++ files and return a [`cc::Build`] you can use to build
    /// more from a build.rs file.
    ///
    /// The error type returned by this function supports [`miette::Diagnostic`],
    /// so if you use the `miette` crate and its `fancy` feature, then simply
    /// return a `miette::Result` from your main function, you should get nicely
    /// printed diagnostics.
    pub fn build(self) -> Result<BuilderBuild, BuilderError> {
        self.build_listing_files().map(|r| r.0)
    }

    /// For use in tests only, this does the build and returns additional information
    /// about the files generated which can subsequently be examined for correctness.
    pub fn build_listing_files(self) -> Result<BuilderSuccess, BuilderError> {
        let clang_args = &self
            .extra_clang_args
            .iter()
            .map(|s| &s[..])
            .collect::<Vec<_>>();
        rust_version_check();
        let gen_location_strategy = match self.custom_gendir {
            None => FileLocationStrategy::new(),
            Some(custom_dir) => FileLocationStrategy::Custom(custom_dir),
        };
        let incdir = gen_location_strategy.get_include_dir();
        ensure_created(&incdir)?;
        let cxxdir = gen_location_strategy.get_cxx_dir();
        ensure_created(&cxxdir)?;
        let rsdir = gen_location_strategy.get_rs_dir();
        ensure_created(&rsdir)?;
        // We are incredibly unsophisticated in our directory arrangement here
        // compared to cxx. I have no doubt that we will need to replicate just
        // about everything cxx does, in due course...
        // Write cxx.h to that location, as it may be needed by
        // some of our generated code.
        write_to_file(
            &incdir,
            "cxx.h",
            &Self::get_cxx_header_bytes(self.cpp_codegen_options.suppress_system_headers),
        )?;

        let autocxx_inc = build_autocxx_inc(self.autocxx_incs, &incdir);
        gen_location_strategy.set_cargo_env_vars_for_build();

        let mut parsed_file = crate::parse_file(self.rs_file, self.auto_allowlist)
            .map_err(BuilderError::ParseError)?;
        parsed_file
            .resolve_all(
                autocxx_inc,
                clang_args,
                self.dependency_recorder,
                &self.cpp_codegen_options,
            )
            .map_err(BuilderError::ParseError)?;
        let mut counter = 0;
        let mut builder = cc::Build::new();
        builder.cpp(true);
        if std::env::var_os("AUTOCXX_ASAN").is_some() {
            builder.flag_if_supported("-fsanitize=address");
        }
        let mut generated_rs = Vec::new();
        let mut generated_cpp = Vec::new();
        builder.includes(parsed_file.include_dirs());
        for include_cpp in parsed_file.get_cpp_buildables() {
            let generated_code = include_cpp
                .generate_h_and_cxx(&self.cpp_codegen_options)
                .map_err(BuilderError::InvalidCxx)?;
            for filepair in generated_code.0 {
                let fname = format!("gen{}.cxx", counter);
                counter += 1;
                if let Some(implementation) = &filepair.implementation {
                    let gen_cxx_path = write_to_file(&cxxdir, &fname, implementation)?;
                    builder.file(&gen_cxx_path);
                    generated_cpp.push(gen_cxx_path);
                }
                write_to_file(&incdir, &filepair.header_name, &filepair.header)?;
                generated_cpp.push(incdir.join(filepair.header_name));
            }
        }

        for include_cpp in parsed_file.get_rs_buildables() {
            let rs = include_cpp.generate_rs();
            generated_rs.push(write_rs_to_file(
                &rsdir,
                &include_cpp.config.get_rs_filename(),
                rs,
            )?);
        }
        if counter == 0 {
            Err(BuilderError::NoIncludeCxxMacrosFound)
        } else {
            Ok(BuilderSuccess(builder, generated_rs, generated_cpp))
        }
    }

    fn get_cxx_header_bytes(suppress_system_headers: bool) -> Vec<u8> {
        strip_system_headers(crate::HEADER.as_bytes().to_vec(), suppress_system_headers)
    }
}

fn ensure_created(dir: &Path) -> Result<(), BuilderError> {
    std::fs::create_dir_all(dir)
        .map_err(|e| BuilderError::UnableToCreateDirectory(e, dir.to_path_buf()))
}

fn build_autocxx_inc<I, T>(paths: I, extra_path: &Path) -> Vec<PathBuf>
where
    I: IntoIterator<Item = T>,
    T: AsRef<OsStr>,
{
    paths
        .into_iter()
        .map(|p| PathBuf::from(p.as_ref()))
        .chain(std::iter::once(extra_path.to_path_buf()))
        .collect()
}

fn write_to_file(dir: &Path, filename: &str, content: &[u8]) -> Result<PathBuf, BuilderError> {
    let path = dir.join(filename);
    try_write_to_file(&path, content).map_err(|e| BuilderError::FileWriteFail(e, path.clone()))?;
    Ok(path)
}

fn try_write_to_file(path: &Path, content: &[u8]) -> std::io::Result<()> {
    let mut f = File::create(path)?;
    f.write_all(content)
}

fn write_rs_to_file(
    dir: &Path,
    filename: &str,
    content: TokenStream,
) -> Result<PathBuf, BuilderError> {
    write_to_file(dir, filename, content.to_string().as_bytes())
}

fn rust_version_check() {
    if !version_check::is_min_version("1.54.0").unwrap_or(false) {
        panic!("Rust 1.54 or later is required.")
    }
}
<commit_msg>Improve Builder documentation.<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use autocxx_parser::file_locations::FileLocationStrategy;
use miette::Diagnostic;
use proc_macro2::TokenStream;
use thiserror::Error;

use crate::{strip_system_headers, CppCodegenOptions, ParseError, RebuildDependencyRecorder};
use std::ffi::OsStr;
use std::ffi::OsString;
use std::fs::File;
use std::io::Write;
use std::marker::PhantomData;
use std::path::{Path, PathBuf};

/// Errors returned during creation of a [`cc::Build`] from an include_cxx
/// macro.
#[derive(Error, Diagnostic, Debug)]
#[cfg_attr(feature = "nightly", doc(cfg(feature = "build")))]
pub enum BuilderError {
    #[error("cxx couldn't handle our generated bindings - could be a bug in autocxx: {0}")]
    InvalidCxx(cxx_gen::Error),
    #[error(transparent)]
    #[diagnostic(transparent)]
    ParseError(ParseError),
    #[error("we couldn't write the generated code to disk at {1}: {0}")]
    FileWriteFail(std::io::Error, PathBuf),
    #[error("no include_cpp! macro was found")]
    NoIncludeCxxMacrosFound,
    #[error("could not create a directory {1}: {0}")]
    UnableToCreateDirectory(std::io::Error, PathBuf),
}

#[cfg_attr(feature = "nightly", doc(cfg(feature = "build")))]
pub type BuilderBuild = cc::Build;

/// For test purposes only, a [`cc::Build`] and lists of Rust and C++
/// files generated.
#[cfg_attr(feature = "nightly", doc(cfg(feature = "build")))]
pub struct BuilderSuccess(pub BuilderBuild, pub Vec<PathBuf>, pub Vec<PathBuf>);

/// Results of a build.
#[cfg_attr(feature = "nightly", doc(cfg(feature = "build")))]
pub type BuilderResult = Result<BuilderSuccess, BuilderError>;

/// The context in which a builder object lives. Callbacks for various
/// purposes.
#[cfg_attr(feature = "nightly", doc(cfg(feature = "build")))]
pub trait BuilderContext {
    /// Perform any initialization specific to the context in which this
    /// builder lives.
    fn setup() {}

    /// Create a dependency recorder, if any.
    fn get_dependency_recorder() -> Option<Box<dyn RebuildDependencyRecorder>>;
}

/// An object to allow building of bindings from a `build.rs` file.
///
/// It would be unusual to use this directly - see the `autocxx_build` or
/// `autocxx_gen` crates.
#[cfg_attr(feature = "nightly", doc(cfg(feature = "build")))]
pub struct Builder<'a, BuilderContext> {
    rs_file: PathBuf,
    autocxx_incs: Vec<OsString>,
    extra_clang_args: Vec<String>,
    dependency_recorder: Option<Box<dyn RebuildDependencyRecorder>>,
    custom_gendir: Option<PathBuf>,
    auto_allowlist: bool,
    cpp_codegen_options: CppCodegenOptions<'a>,
    // This member is to ensure that this type is parameterized
    // by a BuilderContext. The goal is to balance three needs:
    // (1) have most of the functionality over in autocxx_engine,
    // (2) expose this type to users of autocxx_build and to
    //     make it easy for callers simply to call Builder::new,
    // (3) ensure that such a Builder does a few tasks specific to its use
    // in a cargo environment.
    ctx: PhantomData<BuilderContext>,
}

impl<CTX: BuilderContext> Builder<'_, CTX> {
    /// Create a new Builder object. You'll need to pass in the Rust file
    /// which contains the bindings (typically an `include_cpp!` macro
    /// though `autocxx` can also handle manually-crafted `cxx::bridge`
    /// bindings), and a list of include directories which should be searched
    /// by autocxx as it tries to hunt for the include files specified
    /// within the `include_cpp!` macro.
    ///
    /// Usually after this you'd call [`build`].
    pub fn new(
        rs_file: impl AsRef<Path>,
        autocxx_incs: impl IntoIterator<Item = impl AsRef<OsStr>>,
    ) -> Self {
        CTX::setup();
        Self {
            rs_file: rs_file.as_ref().to_path_buf(),
            autocxx_incs: autocxx_incs
                .into_iter()
                .map(|s| s.as_ref().to_os_string())
                .collect(),
            extra_clang_args: Vec::new(),
            dependency_recorder: CTX::get_dependency_recorder(),
            custom_gendir: None,
            auto_allowlist: false,
            cpp_codegen_options: CppCodegenOptions::default(),
            ctx: PhantomData,
        }
    }

    /// Specify extra arguments for clang.
    pub fn extra_clang_args(mut self, extra_clang_args: &[&str]) -> Self {
        self.extra_clang_args = extra_clang_args.iter().map(|s| s.to_string()).collect();
        self
    }

    /// Where to generate the code.
    pub fn custom_gendir(mut self, custom_gendir: PathBuf) -> Self {
        self.custom_gendir = Some(custom_gendir);
        self
    }

    /// Update C++ code generation options. See [`CppCodegenOptions`] for details.
    pub fn cpp_codegen_options<F>(mut self, modifier: F) -> Self
    where
        F: FnOnce(&mut CppCodegenOptions),
    {
        modifier(&mut self.cpp_codegen_options);
        self
    }

    /// Automatically discover uses of the C++ `ffi` mod and generate the allowlist
    /// from that.
    /// This is a highly experimental option, not currently recommended.
    /// It doesn't work in the following cases:
    /// * Static function calls on types within the FFI mod.
    /// * Anything inside a macro invocation.
    /// * You're using a different name for your `ffi` mod
    /// * You're using multiple FFI mods
    /// * You've got usages scattered across files beyond that with the
    ///   `include_cpp` invocation
    /// * You're using `use` statements to rename mods or items. If this
    /// proves to be a promising or helpful direction, autocxx would be happy
    /// to accept pull requests to remove some of these limitations.
    pub fn auto_allowlist(mut self, do_it: bool) -> Self {
        self.auto_allowlist = do_it;
        self
    }

    /// Whether to suppress inclusion of system headers (`memory`, `string` etc.)
    /// from generated C++ bindings code. This should not normally be used,
    /// but can occasionally be useful if you're reducing a test case and you
    /// have a preprocessed header file which already contains absolutely everything
    /// that the bindings could ever need.
    pub fn suppress_system_headers(mut self, do_it: bool) -> Self {
        self.cpp_codegen_options.suppress_system_headers = do_it;
        self
    }

    /// An annotation optionally to include on each C++ function.
    /// For example to export the symbol from a library.
    pub fn cxx_impl_annotations(mut self, cxx_impl_annotations: Option<String>) -> Self {
        self.cpp_codegen_options.cxx_impl_annotations = cxx_impl_annotations;
        self
    }

    /// Whether to skip using [`cxx_gen`] to generate the C++ code,
    /// so that some other process can handle that.
    pub fn skip_cxx_gen(mut self, skip_cxx_gen: bool) -> Self {
        self.cpp_codegen_options.skip_cxx_gen = skip_cxx_gen;
        self
    }

    /// Build autocxx C++ files and return a [`cc::Build`] you can use to build
    /// more from a build.rs file.
    ///
    /// The error type returned by this function supports [`miette::Diagnostic`],
    /// so if you use the `miette` crate and its `fancy` feature, then simply
    /// return a `miette::Result` from your main function, you should get nicely
    /// printed diagnostics.
    pub fn build(self) -> Result<BuilderBuild, BuilderError> {
        self.build_listing_files().map(|r| r.0)
    }

    /// For use in tests only, this does the build and returns additional information
    /// about the files generated which can subsequently be examined for correctness.
    /// In production, please use simply [`build`].
    pub fn build_listing_files(self) -> Result<BuilderSuccess, BuilderError> {
        let clang_args = &self
            .extra_clang_args
            .iter()
            .map(|s| &s[..])
            .collect::<Vec<_>>();
        rust_version_check();
        let gen_location_strategy = match self.custom_gendir {
            None => FileLocationStrategy::new(),
            Some(custom_dir) => FileLocationStrategy::Custom(custom_dir),
        };
        let incdir = gen_location_strategy.get_include_dir();
        ensure_created(&incdir)?;
        let cxxdir = gen_location_strategy.get_cxx_dir();
        ensure_created(&cxxdir)?;
        let rsdir = gen_location_strategy.get_rs_dir();
        ensure_created(&rsdir)?;
        // We are incredibly unsophisticated in our directory arrangement here
        // compared to cxx. I have no doubt that we will need to replicate just
        // about everything cxx does, in due course...
        // Write cxx.h to that location, as it may be needed by
        // some of our generated code.
        write_to_file(
            &incdir,
            "cxx.h",
            &Self::get_cxx_header_bytes(self.cpp_codegen_options.suppress_system_headers),
        )?;

        let autocxx_inc = build_autocxx_inc(self.autocxx_incs, &incdir);
        gen_location_strategy.set_cargo_env_vars_for_build();

        let mut parsed_file = crate::parse_file(self.rs_file, self.auto_allowlist)
            .map_err(BuilderError::ParseError)?;
        parsed_file
            .resolve_all(
                autocxx_inc,
                clang_args,
                self.dependency_recorder,
                &self.cpp_codegen_options,
            )
            .map_err(BuilderError::ParseError)?;
        let mut counter = 0;
        let mut builder = cc::Build::new();
        builder.cpp(true);
        if std::env::var_os("AUTOCXX_ASAN").is_some() {
            builder.flag_if_supported("-fsanitize=address");
        }
        let mut generated_rs = Vec::new();
        let mut generated_cpp = Vec::new();
        builder.includes(parsed_file.include_dirs());
        for include_cpp in parsed_file.get_cpp_buildables() {
            let generated_code = include_cpp
                .generate_h_and_cxx(&self.cpp_codegen_options)
                .map_err(BuilderError::InvalidCxx)?;
            for filepair in generated_code.0 {
                let fname = format!("gen{}.cxx", counter);
                counter += 1;
                if let Some(implementation) = &filepair.implementation {
                    let gen_cxx_path = write_to_file(&cxxdir, &fname, implementation)?;
                    builder.file(&gen_cxx_path);
                    generated_cpp.push(gen_cxx_path);
                }
                write_to_file(&incdir, &filepair.header_name, &filepair.header)?;
                generated_cpp.push(incdir.join(filepair.header_name));
            }
        }

        for include_cpp in parsed_file.get_rs_buildables() {
            let rs = include_cpp.generate_rs();
            generated_rs.push(write_rs_to_file(
                &rsdir,
                &include_cpp.config.get_rs_filename(),
                rs,
            )?);
        }
        if counter == 0 {
            Err(BuilderError::NoIncludeCxxMacrosFound)
        } else {
            Ok(BuilderSuccess(builder, generated_rs, generated_cpp))
        }
    }

    fn get_cxx_header_bytes(suppress_system_headers: bool) -> Vec<u8> {
        strip_system_headers(crate::HEADER.as_bytes().to_vec(), suppress_system_headers)
    }
}

fn ensure_created(dir: &Path) -> Result<(), BuilderError> {
    std::fs::create_dir_all(dir)
        .map_err(|e| BuilderError::UnableToCreateDirectory(e, dir.to_path_buf()))
}

fn build_autocxx_inc<I, T>(paths: I, extra_path: &Path) -> Vec<PathBuf>
where
    I: IntoIterator<Item = T>,
    T: AsRef<OsStr>,
{
    paths
        .into_iter()
        .map(|p| PathBuf::from(p.as_ref()))
        .chain(std::iter::once(extra_path.to_path_buf()))
        .collect()
}

fn write_to_file(dir: &Path, filename: &str, content: &[u8]) -> Result<PathBuf, BuilderError> {
    let path = dir.join(filename);
    try_write_to_file(&path, content).map_err(|e| BuilderError::FileWriteFail(e, path.clone()))?;
    Ok(path)
}

fn try_write_to_file(path: &Path, content: &[u8]) -> std::io::Result<()> {
    let mut f = File::create(path)?;
    f.write_all(content)
}

fn write_rs_to_file(
    dir: &Path,
    filename: &str,
    content: TokenStream,
) -> Result<PathBuf, BuilderError> {
    write_to_file(dir, filename, content.to_string().as_bytes())
}

fn rust_version_check() {
    if !version_check::is_min_version("1.54.0").unwrap_or(false) {
        panic!("Rust 1.54 or later is required.")
    }
}
<|endoftext|>


<commit_before>//< Helpers that are useful for tests and doctests.
use crate::{
    arch::Arch,
    aspace::{AddressSpace, RelativeAddressSpace},
    module::{Module, Permissions, Section},
    RVA, VA,
};

/// configure a global logger at level==DEBUG.
pub fn init_logging() {
    let log_level = log::LevelFilter::Debug;
    fern::Dispatch::new()
        .format(move |out, message, record| {
            out.finish(format_args!(
                "{} [{:5}] {} {}",
                chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
                record.level(),
                if log_level == log::LevelFilter::Trace {
                    record.target()
                } else {
                    ""
                },
                message
            ))
        })
        .level(log_level)
        .chain(std::io::stderr())
        .filter(|metadata| !metadata.target().starts_with("goblin::pe"))
        .apply()
        .expect("failed to configure logging");
}

/// this is for testing, so will panic on error.
pub fn load_shellcode(arch: Arch, buf: &[u8]) -> Module {
    let mut address_space = RelativeAddressSpace::with_capacity(buf.len() as u64);
    address_space.map.writezx(0x0, buf).unwrap();

    Module {
        arch,
        sections: vec![Section {
            name:           "shellcode".to_string(),
            permissions:    Permissions::RWX,
            physical_range: std::ops::Range {
                start: 0x0,
                end:   buf.len() as RVA,
            },
            virtual_range:  std::ops::Range {
                start: 0x0,
                end:   buf.len() as RVA,
            },
        }],
        address_space: address_space.into_absolute(0x0).unwrap(),
    }
}

/// this is for testing, so will panic on error.
pub fn load_shellcode32(buf: &[u8]) -> Module {
    load_shellcode(Arch::X32, buf)
}

/// this is for testing, so will panic on error.
pub fn load_shellcode64(buf: &[u8]) -> Module {
    load_shellcode(Arch::X64, buf)
}

/// this is for testing, so will panic on error.
pub fn read_insn(module: &Module, va: VA) -> zydis::DecodedInstruction {
    use crate::analysis::dis;

    let decoder = dis::get_disassembler(module).unwrap();
    let mut insn_buf = [0u8; 16];
    module.address_space.read_into(va, &mut insn_buf).unwrap();
    decoder.decode(&insn_buf).unwrap().unwrap()
}

pub fn emu_from_shellcode64(code: &[u8]) -> crate::emu::Emulator {
    let m = load_shellcode64(code);
    let mut emu = crate::emu::Emulator::from_module(&m);
    emu.reg.rip = m.address_space.base_address; // 0x0

    emu.mem.mmap(0x5000, 0x2000, Permissions::RW).unwrap();
    emu.reg.rsp = 0x6000;
    emu.reg.rbp = 0x6000;

    emu
}

pub fn emu_from_shellcode32(code: &[u8]) -> crate::emu::Emulator {
    let m = load_shellcode32(code);
    let mut emu = crate::emu::Emulator::from_module(&m);
    emu.reg.set_eip(m.address_space.base_address as u32); // 0x0

    emu.mem.mmap(0x5000, 0x2000, Permissions::RW).unwrap();
    emu.reg.set_esp(0x6000);
    emu.reg.set_ebp(0x6000);

    emu
}

#[cfg(test)]
pub mod uc {
    use byteorder::{ByteOrder, LittleEndian};
    use unicorn::Cpu;

    use super::load_shellcode64;
    use crate::{
        aspace::AddressSpace,
        emu::{mmu::PAGE_SIZE, reg::STATUS_MASK},
        module::{Module, Permissions},
    };

    pub struct Uc {
        emu: unicorn::CpuX86,
    }

    impl Uc {
        pub fn from_module(m: &Module) -> Uc {
            let emu = match m.arch {
                crate::arch::Arch::X32 => unicorn::CpuX86::new(unicorn::Mode::MODE_32).unwrap(),
                crate::arch::Arch::X64 => unicorn::CpuX86::new(unicorn::Mode::MODE_64).unwrap(),
            };

            for section in m.sections.iter() {
                let mut page_addr = section.virtual_range.start;

                let section_size = section.virtual_range.end - section.virtual_range.start;
                emu.mem_map(
                    section.virtual_range.start,
                    crate::util::align(section_size, PAGE_SIZE as u64) as usize,
                    unicorn::PROT_WRITE,
                )
                .unwrap();

                while page_addr < section.virtual_range.end {
                    let mut page = [0u8; PAGE_SIZE];

                    // AddressSpace currently allows non-page-aligned sizes.
                    let page_data = if page_addr + PAGE_SIZE as u64 > section.virtual_range.end {
                        &mut page[..(section.virtual_range.end - page_addr) as usize]
                    } else {
                        &mut page[..]
                    };

                    m.address_space.read_into(page_addr, page_data).unwrap();
                    emu.mem_write(page_addr, &page[..]).unwrap();
                    page_addr += PAGE_SIZE as u64;
                }

                let mut prot = unicorn::PROT_NONE;
                if section.permissions.intersects(Permissions::W) {
                    prot.insert(unicorn::PROT_WRITE);
                }
                if section.permissions.intersects(Permissions::R) {
                    prot.insert(unicorn::PROT_READ);
                }
                if section.permissions.intersects(Permissions::X) {
                    prot.insert(unicorn::PROT_EXEC);
                }

                emu.mem_protect(
                    section.virtual_range.start,
                    crate::util::align(section_size, PAGE_SIZE as u64) as usize,
                    prot,
                )
                .unwrap();
            }

            Uc { emu }
        }

        pub fn step(&mut self) -> Result<(), unicorn::Error> {
            let rip = self.emu.reg_read(unicorn::RegisterX86::RIP).unwrap();
            self.emu.emu_start(rip, u64::MAX, 0, 1)
        }

        pub fn mem_read_u64(&self, addr: u64) -> Result<u64, unicorn::Error> {
            let buf = self.emu.mem_read(addr, 8)?;
            Ok(LittleEndian::read_u64(&buf))
        }

        /// panics if the given emulator does have the same state as this.
        pub fn check(&self, other: &crate::emu::Emulator) {
            use unicorn::RegisterX86::*;

            assert_eq!(self.emu.reg_read(RIP).unwrap(), other.reg.rip(), "register: rip",);

            assert_eq!(self.emu.reg_read(RSP).unwrap(), other.reg.rsp(), "register: rsp",);

            assert_eq!(self.emu.reg_read(RBP).unwrap(), other.reg.rbp(), "register: rbp",);

            assert_eq!(self.emu.reg_read(RAX).unwrap(), other.reg.rax(), "register: rax",);

            assert_eq!(self.emu.reg_read(RBX).unwrap(), other.reg.rbx(), "register: rbx",);

            assert_eq!(self.emu.reg_read(RCX).unwrap(), other.reg.rcx(), "register: rcx",);

            assert_eq!(self.emu.reg_read(RDX).unwrap(), other.reg.rdx(), "register: rdx",);

            assert_eq!(self.emu.reg_read(RSI).unwrap(), other.reg.rsi(), "register: rsi",);

            assert_eq!(self.emu.reg_read(RDI).unwrap(), other.reg.rdi(), "register: rdi",);

            // we don't emulate all of the flags, just the status flags.
            assert_eq!(
                self.emu.reg_read(EFLAGS).unwrap() as u64 & STATUS_MASK,
                other.reg.rflags() & STATUS_MASK,
                "flags"
            );

            // dereferece $SP and $BP
            let rsp = other.reg.rsp();
            assert_eq!(self.mem_read_u64(rsp).unwrap(), other.mem.read_u64(rsp).unwrap(),);

            let rbp = other.reg.rbp();
            assert_eq!(self.mem_read_u64(rbp).unwrap(), other.mem.read_u64(rbp).unwrap(),);
        }
    }

    pub fn uc_from_shellcode64(code: &[u8]) -> Uc {
        let m = load_shellcode64(code);
        let mut uc = Uc::from_module(&m);

        uc.emu
            .reg_write(unicorn::RegisterX86::RIP, m.address_space.base_address)
            .unwrap(); // 0x0

        uc.emu.mem_map(0x5000, 0x2000, unicorn::PROT_ALL).unwrap();
        uc.emu.reg_write(unicorn::RegisterX86::RSP, 0x6000).unwrap();
        uc.emu.reg_write(unicorn::RegisterX86::RBP, 0x6000).unwrap();

        uc
    }
}
<commit_msg>test: emu: better format panic message<commit_after>//< Helpers that are useful for tests and doctests.
use crate::{
    arch::Arch,
    aspace::{AddressSpace, RelativeAddressSpace},
    module::{Module, Permissions, Section},
    RVA, VA,
};

/// configure a global logger at level==DEBUG.
pub fn init_logging() {
    let log_level = log::LevelFilter::Debug;
    fern::Dispatch::new()
        .format(move |out, message, record| {
            out.finish(format_args!(
                "{} [{:5}] {} {}",
                chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
                record.level(),
                if log_level == log::LevelFilter::Trace {
                    record.target()
                } else {
                    ""
                },
                message
            ))
        })
        .level(log_level)
        .chain(std::io::stderr())
        .filter(|metadata| !metadata.target().starts_with("goblin::pe"))
        .apply()
        .expect("failed to configure logging");
}

/// this is for testing, so will panic on error.
pub fn load_shellcode(arch: Arch, buf: &[u8]) -> Module {
    let mut address_space = RelativeAddressSpace::with_capacity(buf.len() as u64);
    address_space.map.writezx(0x0, buf).unwrap();

    Module {
        arch,
        sections: vec![Section {
            name:           "shellcode".to_string(),
            permissions:    Permissions::RWX,
            physical_range: std::ops::Range {
                start: 0x0,
                end:   buf.len() as RVA,
            },
            virtual_range:  std::ops::Range {
                start: 0x0,
                end:   buf.len() as RVA,
            },
        }],
        address_space: address_space.into_absolute(0x0).unwrap(),
    }
}

/// this is for testing, so will panic on error.
pub fn load_shellcode32(buf: &[u8]) -> Module {
    load_shellcode(Arch::X32, buf)
}

/// this is for testing, so will panic on error.
pub fn load_shellcode64(buf: &[u8]) -> Module {
    load_shellcode(Arch::X64, buf)
}

/// this is for testing, so will panic on error.
pub fn read_insn(module: &Module, va: VA) -> zydis::DecodedInstruction {
    use crate::analysis::dis;

    let decoder = dis::get_disassembler(module).unwrap();
    let mut insn_buf = [0u8; 16];
    module.address_space.read_into(va, &mut insn_buf).unwrap();
    decoder.decode(&insn_buf).unwrap().unwrap()
}

pub fn emu_from_shellcode64(code: &[u8]) -> crate::emu::Emulator {
    let m = load_shellcode64(code);
    let mut emu = crate::emu::Emulator::from_module(&m);
    emu.reg.rip = m.address_space.base_address; // 0x0

    emu.mem.mmap(0x5000, 0x2000, Permissions::RW).unwrap();
    emu.reg.rsp = 0x6000;
    emu.reg.rbp = 0x6000;

    emu
}

pub fn emu_from_shellcode32(code: &[u8]) -> crate::emu::Emulator {
    let m = load_shellcode32(code);
    let mut emu = crate::emu::Emulator::from_module(&m);
    emu.reg.set_eip(m.address_space.base_address as u32); // 0x0

    emu.mem.mmap(0x5000, 0x2000, Permissions::RW).unwrap();
    emu.reg.set_esp(0x6000);
    emu.reg.set_ebp(0x6000);

    emu
}

#[cfg(test)]
pub mod uc {
    use byteorder::{ByteOrder, LittleEndian};
    use unicorn::Cpu;

    use super::load_shellcode64;
    use crate::{
        aspace::AddressSpace,
        emu::{mmu::PAGE_SIZE, reg::STATUS_MASK},
        module::{Module, Permissions},
    };

    pub struct Uc {
        emu: unicorn::CpuX86,
    }

    impl Uc {
        pub fn from_module(m: &Module) -> Uc {
            let emu = match m.arch {
                crate::arch::Arch::X32 => unicorn::CpuX86::new(unicorn::Mode::MODE_32).unwrap(),
                crate::arch::Arch::X64 => unicorn::CpuX86::new(unicorn::Mode::MODE_64).unwrap(),
            };

            for section in m.sections.iter() {
                let mut page_addr = section.virtual_range.start;

                let section_size = section.virtual_range.end - section.virtual_range.start;
                emu.mem_map(
                    section.virtual_range.start,
                    crate::util::align(section_size, PAGE_SIZE as u64) as usize,
                    unicorn::PROT_WRITE,
                )
                .unwrap();

                while page_addr < section.virtual_range.end {
                    let mut page = [0u8; PAGE_SIZE];

                    // AddressSpace currently allows non-page-aligned sizes.
                    let page_data = if page_addr + PAGE_SIZE as u64 > section.virtual_range.end {
                        &mut page[..(section.virtual_range.end - page_addr) as usize]
                    } else {
                        &mut page[..]
                    };

                    m.address_space.read_into(page_addr, page_data).unwrap();
                    emu.mem_write(page_addr, &page[..]).unwrap();
                    page_addr += PAGE_SIZE as u64;
                }

                let mut prot = unicorn::PROT_NONE;
                if section.permissions.intersects(Permissions::W) {
                    prot.insert(unicorn::PROT_WRITE);
                }
                if section.permissions.intersects(Permissions::R) {
                    prot.insert(unicorn::PROT_READ);
                }
                if section.permissions.intersects(Permissions::X) {
                    prot.insert(unicorn::PROT_EXEC);
                }

                emu.mem_protect(
                    section.virtual_range.start,
                    crate::util::align(section_size, PAGE_SIZE as u64) as usize,
                    prot,
                )
                .unwrap();
            }

            Uc { emu }
        }

        pub fn step(&mut self) -> Result<(), unicorn::Error> {
            let rip = self.emu.reg_read(unicorn::RegisterX86::RIP).unwrap();
            self.emu.emu_start(rip, u64::MAX, 0, 1)
        }

        pub fn mem_read_u64(&self, addr: u64) -> Result<u64, unicorn::Error> {
            let buf = self.emu.mem_read(addr, 8)?;
            Ok(LittleEndian::read_u64(&buf))
        }

        /// panics if the given emulator does have the same state as this.
        pub fn check(&self, other: &crate::emu::Emulator) {
            use unicorn::RegisterX86::*;

            assert_eq!(
                self.emu.reg_read(RIP).unwrap(),
                other.reg.rip(),
                "register: rip, uc: {:#x} emu: {:#x}",
                self.emu.reg_read(RIP).unwrap(),
                other.reg.rip()
            );

            assert_eq!(
                self.emu.reg_read(RSP).unwrap(),
                other.reg.rsp(),
                "register: rsp, uc: {:#x} emu: {:#x}",
                self.emu.reg_read(RSP).unwrap(),
                other.reg.rsp()
            );

            assert_eq!(
                self.emu.reg_read(RBP).unwrap(),
                other.reg.rbp(),
                "register: rbp, uc: {:#x} emu: {:#x}",
                self.emu.reg_read(RBP).unwrap(),
                other.reg.rbp()
            );

            assert_eq!(
                self.emu.reg_read(RAX).unwrap(),
                other.reg.rax(),
                "register: rax, uc: {:#x} emu: {:#x}",
                self.emu.reg_read(RAX).unwrap(),
                other.reg.rax()
            );

            assert_eq!(
                self.emu.reg_read(RBX).unwrap(),
                other.reg.rbx(),
                "register: rbx, uc: {:#x} emu: {:#x}",
                self.emu.reg_read(RBX).unwrap(),
                other.reg.rbx()
            );

            assert_eq!(
                self.emu.reg_read(RCX).unwrap(),
                other.reg.rcx(),
                "register: rcx, uc: {:#x} emu: {:#x}",
                self.emu.reg_read(RCX).unwrap(),
                other.reg.rcx()
            );

            assert_eq!(
                self.emu.reg_read(RDX).unwrap(),
                other.reg.rdx(),
                "register: rdx, uc: {:#x} emu: {:#x}",
                self.emu.reg_read(RDX).unwrap(),
                other.reg.rdx()
            );

            assert_eq!(
                self.emu.reg_read(RSI).unwrap(),
                other.reg.rsi(),
                "register: rsi, uc: {:#x} emu: {:#x}",
                self.emu.reg_read(RSI).unwrap(),
                other.reg.rsi()
            );

            assert_eq!(
                self.emu.reg_read(RDI).unwrap(),
                other.reg.rdi(),
                "register: rdi, uc: {:#x} emu: {:#x}",
                self.emu.reg_read(RDI).unwrap(),
                other.reg.rdi()
            );

            // we don't emulate all of the flags, just the status flags.
            assert_eq!(
                self.emu.reg_read(EFLAGS).unwrap() as u64 & STATUS_MASK,
                other.reg.rflags() & STATUS_MASK,
                "flags, uc: {:#x} emu: {:#x}",
                self.emu.reg_read(EFLAGS).unwrap() as u64 & STATUS_MASK,
                other.reg.rflags() & STATUS_MASK,
            );

            // dereferece $SP and $BP
            let rsp = other.reg.rsp();
            assert_eq!(
                self.mem_read_u64(rsp).unwrap(),
                other.mem.read_u64(rsp).unwrap(),
                "*RSP, uc: {:#x} emu: {:#x}",
                self.mem_read_u64(rsp).unwrap(),
                other.mem.read_u64(rsp).unwrap(),
            );

            let rbp = other.reg.rbp();
            assert_eq!(
                self.mem_read_u64(rbp).unwrap(),
                other.mem.read_u64(rbp).unwrap(),
                "*RBP, uc: {:#x} emu: {:#x}",
                self.mem_read_u64(rbp).unwrap(),
                other.mem.read_u64(rbp).unwrap(),
            );
        }
    }

    pub fn uc_from_shellcode64(code: &[u8]) -> Uc {
        let m = load_shellcode64(code);
        let mut uc = Uc::from_module(&m);

        uc.emu
            .reg_write(unicorn::RegisterX86::RIP, m.address_space.base_address)
            .unwrap(); // 0x0

        uc.emu.mem_map(0x5000, 0x2000, unicorn::PROT_ALL).unwrap();
        uc.emu.reg_write(unicorn::RegisterX86::RSP, 0x6000).unwrap();
        uc.emu.reg_write(unicorn::RegisterX86::RBP, 0x6000).unwrap();

        uc
    }
}
<|endoftext|>


<commit_before>extern crate libc;
extern crate sys_util;
extern crate kvm_sys;
extern crate kvm;
extern crate kernel_loader;
extern crate x86_64;
extern crate clap;
extern crate devices;

pub mod machine;

use std::ffi::{CStr, CString};
use std::fs::File;
use std::io::{self, stdout, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Barrier, Mutex};
use std::thread;
use std::thread::{JoinHandle};
use kvm::*;
use kvm_sys::kvm_regs;
use sys_util::{register_signal_handler, EventFd, GuestAddress, GuestMemory,
               Killable, Pollable, Poller, Terminal};
use machine::MachineCfg;

const KERNEL_START_OFFSET: usize = 0x200000;
const CMDLINE_OFFSET: usize = 0x20000;

#[derive(Debug)]
pub enum Error {
    ConfigureSystem(x86_64::Error),
    EventFd(sys_util::Error),
    GuestMemory(sys_util::GuestMemoryError),
    Irq(sys_util::Error),
    Kernel(std::io::Error),
    KernelLoader(kernel_loader::Error),
    Kvm(sys_util::Error),
    Vcpu(sys_util::Error),
    VcpuSpawn(std::io::Error),
    Vm(sys_util::Error),
}

impl std::convert::From<kernel_loader::Error> for Error {
    fn from(e: kernel_loader::Error) -> Error {
        Error::KernelLoader(e)
    }
}

impl std::convert::From<x86_64::Error> for Error {
    fn from(e: x86_64::Error) -> Error {
        Error::ConfigureSystem(e)
    }
}

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

pub fn boot_kernel(cfg: &MachineCfg) -> Result<()> {
    // FIXME branciog@ do not hardcode the vm mem size
    // Hardcoding the vm memory size to 128MB
    let mem_size = 128 << 20;
    let arch_mem_regions = x86_64::arch_memory_regions(mem_size);

    let mut kernel_file;
    match cfg.kernel_path {
        Some(ref kernel_path) => {
            kernel_file = File::open(kernel_path.as_path())
                    .map_err(Error::Kernel)?
        },
        None => {
            return Err(Error::Kernel(
                    io::Error::new(io::ErrorKind::NotFound,
                                   "missing kernel path")))
        }
    }

    let cmdline: CString = match cfg.kernel_cmdline {
        Some(ref v) => CString::new(v.as_bytes()).unwrap(),
        _ => return Err(Error::Kernel(
                io::Error::new(io::ErrorKind::NotFound,
                               "missing kernel cmdline")))
    };
    let cmdline: &CStr = &cmdline;
    let vcpu_count = 1;

    let kernel_start_addr = GuestAddress(KERNEL_START_OFFSET);
    let cmdline_addr = GuestAddress(CMDLINE_OFFSET);

    let guest_mem = GuestMemory::new(&arch_mem_regions)
            .map_err(Error::GuestMemory)?;

    let kvm = Kvm::new().map_err(Error::Kvm)?;
    let vm = Vm::new(&kvm, guest_mem).map_err(Error::Vm)?;

    let tss_addr = GuestAddress(0xfffbd000);
    vm.set_tss_addr(tss_addr).map_err(Error::Vm)?;
    vm.create_pit().map_err(Error::Vm)?;
    vm.create_irq_chip().map_err(Error::Vm)?;

    kernel_loader::load_kernel(vm.get_memory(), kernel_start_addr,
                               &mut kernel_file)?;
    kernel_loader::load_cmdline(vm.get_memory(), cmdline_addr, cmdline)?;

    x86_64::configure_system(vm.get_memory(),
                             kernel_start_addr,
                             cmdline_addr,
                             cmdline.to_bytes().len() + 1,
                             vcpu_count as u8)?;

    let mut io_bus = devices::Bus::new();
    let com_evt_1_3 = EventFd::new().map_err(Error::EventFd)?;
    let com_evt_2_4 = EventFd::new().map_err(Error::EventFd)?;
    let stdio_serial = Arc::new(Mutex::new(devices::Serial::new_out(
            com_evt_1_3.try_clone().map_err(Error::EventFd)?,
            Box::new(stdout()))));

    io_bus.insert(stdio_serial.clone(), 0x3f8, 0x8).unwrap();
    io_bus.insert(Arc::new(Mutex::new(devices::Serial::new_sink(
            com_evt_2_4.try_clone().map_err(Error::EventFd)?))), 0x2f8, 0x8)
            .unwrap();
    io_bus.insert(Arc::new(Mutex::new(devices::Serial::new_sink(
            com_evt_1_3.try_clone().map_err(Error::EventFd)?))), 0x3e8, 0x8)
            .unwrap();
    io_bus.insert(Arc::new(Mutex::new(devices::Serial::new_sink(
            com_evt_2_4.try_clone().map_err(Error::EventFd)?))), 0x2e8, 0x8)
            .unwrap();

    vm.register_irqfd(&com_evt_1_3, 4)
            .map_err(Error::Irq)?;
    vm.register_irqfd(&com_evt_2_4, 3)
            .map_err(Error::Irq)?;

    let exit_evt = EventFd::new().map_err(Error::EventFd)?;
    let mut vcpu_handles = Vec::with_capacity(vcpu_count as usize);
    let vcpu_thread_barrier = Arc::new(Barrier::new((vcpu_count + 1) as usize));
    let kill_signaled = Arc::new(AtomicBool::new(false));

    for cpu_id in 0..vcpu_count {
        let io_bus = io_bus.clone();
        let kill_signaled = kill_signaled.clone();
        let vcpu_thread_barrier = vcpu_thread_barrier.clone();
        let vcpu_exit_evt = exit_evt.try_clone().map_err(Error::EventFd)?;

        let vcpu = Vcpu::new(cpu_id as libc::c_ulong, &kvm, &vm).map_err(Error::Vcpu)?;
        x86_64::configure_vcpu(vm.get_memory(), kernel_start_addr, &kvm,
                               &vcpu, cpu_id as u64, vcpu_count as u64)?;
        vcpu_handles.push(thread::Builder::new()
                .name(format!("fc_vcpu{}", cpu_id))
                .spawn(move || {
            unsafe {
                extern "C" fn handle_signal() {}
                // Our signal handler does nothing and is trivially async signal safe.
                register_signal_handler(0, handle_signal)
                        .expect("failed to register vcpu signal handler");
            }

            vcpu_thread_barrier.wait();

            loop {
                match vcpu.run() {
                    Ok(run) => match run {
                        VcpuExit::IoIn(addr, data) => {
                            io_bus.read(addr as u64, data);
                        },
                        VcpuExit::IoOut(addr, data) => {
                            io_bus.write(addr as u64, data);
                        },
                        VcpuExit::MmioRead(_, _) => {},
                        VcpuExit::MmioWrite(_, _) => {},
                        VcpuExit::Hlt => {
                            println!("KVM_EXIT_HLT");
                            break;
                        },
                        VcpuExit::Shutdown => {
                            println!("KVM_EXIT_SHUTDOWN");
                            break;
                        },
                        r => {
                            println!("unexpected exit reason: {:?}", r);
                            break;
                        }
                    },
                    Err(e) => {
                        match e.errno() {
                            libc::EAGAIN | libc::EINTR => {},
                            _ => {
                                println!("vcpu hit unknown error: {:?}", e);
                                break;
                            }
                        }
                    }
                }

                if kill_signaled.load(Ordering::SeqCst) {
                    break;
                }
            }

            vcpu_exit_evt
                .write(1)
                .expect("failed to signal vcpu exit eventfd");

        }).map_err(Error::VcpuSpawn)?);
    }

    vcpu_thread_barrier.wait();

    run_control(stdio_serial, exit_evt, kill_signaled, vcpu_handles)
}

fn run_control(stdio_serial: Arc<Mutex<devices::Serial>>,
               exit_evt: EventFd,
               kill_signaled: Arc<AtomicBool>,
               vcpu_handles: Vec<JoinHandle<()>>) -> Result<()> {
    const EXIT: u32 = 0;
    const STDIN: u32 = 1;

    let stdin_handle = io::stdin();
    let stdin_lock = stdin_handle.lock();
    stdin_lock
            .set_raw_mode()
            .expect("failed to set terminal raw mode");

    let mut pollables = Vec::new();
    pollables.push((EXIT, &exit_evt as &Pollable));
    pollables.push((STDIN, &stdin_lock as &Pollable));

    let mut poller = Poller::new(pollables.len());

    'poll: loop {
        let tokens = match poller.poll(&pollables[..]) {
            Ok(v) => v,
            Err(e) => {
                println!("failed to poll: {:?}", e);
                break;
            }
        };

        for &token in tokens {
            match token {
                EXIT => {
                    println!("vcpu requested shutdown");
                    break 'poll;
                },
                STDIN => {
                    let mut out = [0u8; 64];
                    match stdin_lock.read_raw(&mut out[..]) {
                        Ok(0) => {
                            // Zero-length read indicates EOF. Remove from pollables.
                            pollables.retain(|&pollable| pollable.0 != STDIN);
                        },
                        Err(e) => {
                            println!("error while reading stdin: {:?}", e);
                            pollables.retain(|&pollable| pollable.0 != STDIN);
                        },
                        Ok(count) => {
                            stdio_serial
                                    .lock()
                                    .unwrap()
                                    .queue_input_bytes(&out[..count])
                                    .expect("failed to queue bytes into serial port");
                        },
                    }
                }
                _ => {}
            }
        }
    }

    kill_signaled.store(true, Ordering::SeqCst);
    for handle in vcpu_handles {
        match handle.kill(0) {
            Ok(_) => {
                if let Err(e) = handle.join() {
                    println!("failed to join vcpu thread: {:?}", e);
                }
            },
            Err(e) => println!("failed to kill vcpu thread: {:?}", e),
        }
    }

    stdin_lock
            .set_canon_mode()
            .expect("failed to restore canonical mode for terminal");

    Ok(())
}

pub fn run_x86_code() {
    // This example based on https://lwn.net/Articles/658511/
    let code = [
        /* mov $0x3f8, %dx */
        0xba,
        0xf8,
        0x03,
        /* add %bl, %al */
        0x00,
        0xd8,
        /* add $'0', %al */
        0x04,
        '0' as u8,
        /* out %al, (%dx) */
        0xee,
        /* mov $'\n', %al */
        0xb0,
        '\n' as u8,
        /* out %al, (%dx) */
        0xee,
        /* hlt */
        0xf4,
    ];

    let mem_size = 0x1000;
    let load_addr = GuestAddress(0x1000);
    let mem = GuestMemory::new(&vec![(load_addr, mem_size)]).unwrap();

    let kvm = Kvm::new().expect("new kvm failed");
    let vm = Vm::new(&kvm, mem).expect("new vm failed");
    let vcpu = Vcpu::new(0, &kvm, &vm).expect("new vcpu failed");

    vm.get_memory()
        .write_slice_at_addr(&code, load_addr)
        .expect("Writing code to memory failed.");

    let mut vcpu_sregs = vcpu.get_sregs().expect("get sregs failed");
    assert_ne!(vcpu_sregs.cs.base, 0);
    assert_ne!(vcpu_sregs.cs.selector, 0);
    vcpu_sregs.cs.base = 0;
    vcpu_sregs.cs.selector = 0;
    vcpu.set_sregs(&vcpu_sregs).expect("set sregs failed");

    let mut vcpu_regs: kvm_regs = unsafe { std::mem::zeroed() };
    vcpu_regs.rip = 0x1000;
    vcpu_regs.rax = 2;
    vcpu_regs.rbx = 2;
    vcpu_regs.rflags = 2;
    vcpu.set_regs(&vcpu_regs).expect("set regs failed");

    loop {
        match vcpu.run().expect("run failed") {
            VcpuExit::IoOut(0x3f8, data) => {
                assert_eq!(data.len(), 1);
                io::stdout().write(data).unwrap();
            },
            VcpuExit::Hlt => {
                io::stdout().write(b"KVM_EXIT_HLT\n").unwrap();
                break
            },
            r => panic!("unexpected exit reason: {:?}", r),
        }
    }
}

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

    #[test]
    fn run_code() {
        run_x86_code();
    }
}
<commit_msg>vmm: fix style by running cargo fmt<commit_after>extern crate libc;
extern crate sys_util;
extern crate kvm_sys;
extern crate kvm;
extern crate kernel_loader;
extern crate x86_64;
extern crate clap;
extern crate devices;

pub mod machine;

use std::ffi::{CStr, CString};
use std::fs::File;
use std::io::{self, stdout, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Barrier, Mutex};
use std::thread;
use std::thread::JoinHandle;
use kvm::*;
use kvm_sys::kvm_regs;
use sys_util::{register_signal_handler, EventFd, GuestAddress, GuestMemory, Killable, Pollable,
               Poller, Terminal};
use machine::MachineCfg;

const KERNEL_START_OFFSET: usize = 0x200000;
const CMDLINE_OFFSET: usize = 0x20000;

#[derive(Debug)]
pub enum Error {
    ConfigureSystem(x86_64::Error),
    EventFd(sys_util::Error),
    GuestMemory(sys_util::GuestMemoryError),
    Irq(sys_util::Error),
    Kernel(std::io::Error),
    KernelLoader(kernel_loader::Error),
    Kvm(sys_util::Error),
    Vcpu(sys_util::Error),
    VcpuSpawn(std::io::Error),
    Vm(sys_util::Error),
}

impl std::convert::From<kernel_loader::Error> for Error {
    fn from(e: kernel_loader::Error) -> Error {
        Error::KernelLoader(e)
    }
}

impl std::convert::From<x86_64::Error> for Error {
    fn from(e: x86_64::Error) -> Error {
        Error::ConfigureSystem(e)
    }
}

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

pub fn boot_kernel(cfg: &MachineCfg) -> Result<()> {
    // FIXME branciog@ do not hardcode the vm mem size
    // Hardcoding the vm memory size to 128MB
    let mem_size = 128 << 20;
    let arch_mem_regions = x86_64::arch_memory_regions(mem_size);

    let mut kernel_file;
    match cfg.kernel_path {
        Some(ref kernel_path) => {
            kernel_file = File::open(kernel_path.as_path()).map_err(Error::Kernel)?
        }
        None => {
            return Err(Error::Kernel(io::Error::new(
                io::ErrorKind::NotFound,
                "missing kernel path",
            )))
        }
    }

    let cmdline: CString = match cfg.kernel_cmdline {
        Some(ref v) => CString::new(v.as_bytes()).unwrap(),
        _ => {
            return Err(Error::Kernel(io::Error::new(
                io::ErrorKind::NotFound,
                "missing kernel cmdline",
            )))
        }
    };
    let cmdline: &CStr = &cmdline;
    let vcpu_count = 1;

    let kernel_start_addr = GuestAddress(KERNEL_START_OFFSET);
    let cmdline_addr = GuestAddress(CMDLINE_OFFSET);

    let guest_mem = GuestMemory::new(&arch_mem_regions).map_err(
        Error::GuestMemory,
    )?;

    let kvm = Kvm::new().map_err(Error::Kvm)?;
    let vm = Vm::new(&kvm, guest_mem).map_err(Error::Vm)?;

    let tss_addr = GuestAddress(0xfffbd000);
    vm.set_tss_addr(tss_addr).map_err(Error::Vm)?;
    vm.create_pit().map_err(Error::Vm)?;
    vm.create_irq_chip().map_err(Error::Vm)?;

    kernel_loader::load_kernel(vm.get_memory(), kernel_start_addr, &mut kernel_file)?;
    kernel_loader::load_cmdline(vm.get_memory(), cmdline_addr, cmdline)?;

    x86_64::configure_system(
        vm.get_memory(),
        kernel_start_addr,
        cmdline_addr,
        cmdline.to_bytes().len() + 1,
        vcpu_count as u8,
    )?;

    let mut io_bus = devices::Bus::new();
    let com_evt_1_3 = EventFd::new().map_err(Error::EventFd)?;
    let com_evt_2_4 = EventFd::new().map_err(Error::EventFd)?;
    let stdio_serial = Arc::new(Mutex::new(devices::Serial::new_out(
        com_evt_1_3.try_clone().map_err(Error::EventFd)?,
        Box::new(stdout()),
    )));

    io_bus.insert(stdio_serial.clone(), 0x3f8, 0x8).unwrap();
    io_bus
        .insert(
            Arc::new(Mutex::new(devices::Serial::new_sink(
                com_evt_2_4.try_clone().map_err(Error::EventFd)?,
            ))),
            0x2f8,
            0x8,
        )
        .unwrap();
    io_bus
        .insert(
            Arc::new(Mutex::new(devices::Serial::new_sink(
                com_evt_1_3.try_clone().map_err(Error::EventFd)?,
            ))),
            0x3e8,
            0x8,
        )
        .unwrap();
    io_bus
        .insert(
            Arc::new(Mutex::new(devices::Serial::new_sink(
                com_evt_2_4.try_clone().map_err(Error::EventFd)?,
            ))),
            0x2e8,
            0x8,
        )
        .unwrap();

    vm.register_irqfd(&com_evt_1_3, 4).map_err(Error::Irq)?;
    vm.register_irqfd(&com_evt_2_4, 3).map_err(Error::Irq)?;

    let exit_evt = EventFd::new().map_err(Error::EventFd)?;
    let mut vcpu_handles = Vec::with_capacity(vcpu_count as usize);
    let vcpu_thread_barrier = Arc::new(Barrier::new((vcpu_count + 1) as usize));
    let kill_signaled = Arc::new(AtomicBool::new(false));

    for cpu_id in 0..vcpu_count {
        let io_bus = io_bus.clone();
        let kill_signaled = kill_signaled.clone();
        let vcpu_thread_barrier = vcpu_thread_barrier.clone();
        let vcpu_exit_evt = exit_evt.try_clone().map_err(Error::EventFd)?;

        let vcpu = Vcpu::new(cpu_id as libc::c_ulong, &kvm, &vm).map_err(
            Error::Vcpu,
        )?;
        x86_64::configure_vcpu(
            vm.get_memory(),
            kernel_start_addr,
            &kvm,
            &vcpu,
            cpu_id as u64,
            vcpu_count as u64,
        )?;
        vcpu_handles.push(thread::Builder::new()
            .name(format!("fc_vcpu{}", cpu_id))
            .spawn(move || {
                unsafe {
                    extern "C" fn handle_signal() {}
                    // Our signal handler does nothing and is trivially async signal safe.
                    register_signal_handler(0, handle_signal).expect(
                        "failed to register vcpu signal handler",
                    );
                }

                vcpu_thread_barrier.wait();

                loop {
                    match vcpu.run() {
                        Ok(run) => {
                            match run {
                                VcpuExit::IoIn(addr, data) => {
                                    io_bus.read(addr as u64, data);
                                }
                                VcpuExit::IoOut(addr, data) => {
                                    io_bus.write(addr as u64, data);
                                }
                                VcpuExit::MmioRead(_, _) => {}
                                VcpuExit::MmioWrite(_, _) => {}
                                VcpuExit::Hlt => {
                                    println!("KVM_EXIT_HLT");
                                    break;
                                }
                                VcpuExit::Shutdown => {
                                    println!("KVM_EXIT_SHUTDOWN");
                                    break;
                                }
                                r => {
                                    println!("unexpected exit reason: {:?}", r);
                                    break;
                                }
                            }
                        }
                        Err(e) => {
                            match e.errno() {
                                libc::EAGAIN | libc::EINTR => {}
                                _ => {
                                    println!("vcpu hit unknown error: {:?}", e);
                                    break;
                                }
                            }
                        }
                    }

                    if kill_signaled.load(Ordering::SeqCst) {
                        break;
                    }
                }

                vcpu_exit_evt.write(1).expect(
                    "failed to signal vcpu exit eventfd",
                );

            })
            .map_err(Error::VcpuSpawn)?);
    }

    vcpu_thread_barrier.wait();

    run_control(stdio_serial, exit_evt, kill_signaled, vcpu_handles)
}

fn run_control(
    stdio_serial: Arc<Mutex<devices::Serial>>,
    exit_evt: EventFd,
    kill_signaled: Arc<AtomicBool>,
    vcpu_handles: Vec<JoinHandle<()>>,
) -> Result<()> {
    const EXIT: u32 = 0;
    const STDIN: u32 = 1;

    let stdin_handle = io::stdin();
    let stdin_lock = stdin_handle.lock();
    stdin_lock.set_raw_mode().expect(
        "failed to set terminal raw mode",
    );

    let mut pollables = Vec::new();
    pollables.push((EXIT, &exit_evt as &Pollable));
    pollables.push((STDIN, &stdin_lock as &Pollable));

    let mut poller = Poller::new(pollables.len());

    'poll: loop {
        let tokens = match poller.poll(&pollables[..]) {
            Ok(v) => v,
            Err(e) => {
                println!("failed to poll: {:?}", e);
                break;
            }
        };

        for &token in tokens {
            match token {
                EXIT => {
                    println!("vcpu requested shutdown");
                    break 'poll;
                }
                STDIN => {
                    let mut out = [0u8; 64];
                    match stdin_lock.read_raw(&mut out[..]) {
                        Ok(0) => {
                            // Zero-length read indicates EOF. Remove from pollables.
                            pollables.retain(|&pollable| pollable.0 != STDIN);
                        }
                        Err(e) => {
                            println!("error while reading stdin: {:?}", e);
                            pollables.retain(|&pollable| pollable.0 != STDIN);
                        }
                        Ok(count) => {
                            stdio_serial
                                .lock()
                                .unwrap()
                                .queue_input_bytes(&out[..count])
                                .expect("failed to queue bytes into serial port");
                        }
                    }
                }
                _ => {}
            }
        }
    }

    kill_signaled.store(true, Ordering::SeqCst);
    for handle in vcpu_handles {
        match handle.kill(0) {
            Ok(_) => {
                if let Err(e) = handle.join() {
                    println!("failed to join vcpu thread: {:?}", e);
                }
            }
            Err(e) => println!("failed to kill vcpu thread: {:?}", e),
        }
    }

    stdin_lock.set_canon_mode().expect(
        "failed to restore canonical mode for terminal",
    );

    Ok(())
}

pub fn run_x86_code() {
    // This example based on https://lwn.net/Articles/658511/
    let code = [
        /* mov $0x3f8, %dx */
        0xba,
        0xf8,
        0x03,
        /* add %bl, %al */
        0x00,
        0xd8,
        /* add $'0', %al */
        0x04,
        '0' as u8,
        /* out %al, (%dx) */
        0xee,
        /* mov $'\n', %al */
        0xb0,
        '\n' as u8,
        /* out %al, (%dx) */
        0xee,
        /* hlt */
        0xf4,
    ];

    let mem_size = 0x1000;
    let load_addr = GuestAddress(0x1000);
    let mem = GuestMemory::new(&vec![(load_addr, mem_size)]).unwrap();

    let kvm = Kvm::new().expect("new kvm failed");
    let vm = Vm::new(&kvm, mem).expect("new vm failed");
    let vcpu = Vcpu::new(0, &kvm, &vm).expect("new vcpu failed");

    vm.get_memory()
        .write_slice_at_addr(&code, load_addr)
        .expect("Writing code to memory failed.");

    let mut vcpu_sregs = vcpu.get_sregs().expect("get sregs failed");
    assert_ne!(vcpu_sregs.cs.base, 0);
    assert_ne!(vcpu_sregs.cs.selector, 0);
    vcpu_sregs.cs.base = 0;
    vcpu_sregs.cs.selector = 0;
    vcpu.set_sregs(&vcpu_sregs).expect("set sregs failed");

    let mut vcpu_regs: kvm_regs = unsafe { std::mem::zeroed() };
    vcpu_regs.rip = 0x1000;
    vcpu_regs.rax = 2;
    vcpu_regs.rbx = 2;
    vcpu_regs.rflags = 2;
    vcpu.set_regs(&vcpu_regs).expect("set regs failed");

    loop {
        match vcpu.run().expect("run failed") {
            VcpuExit::IoOut(0x3f8, data) => {
                assert_eq!(data.len(), 1);
                io::stdout().write(data).unwrap();
            }
            VcpuExit::Hlt => {
                io::stdout().write(b"KVM_EXIT_HLT\n").unwrap();
                break;
            }
            r => panic!("unexpected exit reason: {:?}", r),
        }
    }
}

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

    #[test]
    fn run_code() {
        run_x86_code();
    }
}
<|endoftext|>


<commit_before>#![cfg(feature = "fms")]

extern crate rusoto_core;
extern crate rusoto_fms;

use rusoto_fms::{Fms, FmsClient, ListPoliciesRequest};
use rusoto_core::Region;

#[test]
fn should_list_policies() {
    let client = FmsClient::new(Region::UsEast1);
    let request = ListPoliciesRequest::default();

    // If our account doesn't have access, assume everything is fine:
    match client.list_policies(request).sync() {
        Err(e) => assert!(format!("{}", e).contains("is not currently delegated by AWS FM")),
        Ok(res) => println!("Got these policies: {:?}", res),
    }

}
<commit_msg>Fix FMS integration test.<commit_after>#![cfg(feature = "fms")]

extern crate rusoto_core;
extern crate rusoto_fms;

use rusoto_fms::{Fms, FmsClient, ListPoliciesRequest, ListPoliciesError};
use rusoto_core::Region;

use std::str;

#[test]
fn should_list_policies() {
    let client = FmsClient::new(Region::UsEast1);
    let request = ListPoliciesRequest::default();

    // If our account doesn't have access, assume everything is fine:
    match client.list_policies(request).sync() {
        Err(e) => {
            match e {
                ListPoliciesError::Unknown(ref e) => assert!(str::from_utf8(&e.body).unwrap().contains("is not currently delegated by AWS FM"), "Missing error message"),
                _ => panic!("Should have a typed error from FMS"),
            }
        },
        Ok(res) => println!("Got these policies: {:?}", res),
    }

}
<|endoftext|>