<commit_before>extern crate glium;
extern crate clap;
extern crate image;
extern crate cgmath;
extern crate time;

use glium::backend::glutin_backend::GlutinFacade;
use clap::ArgMatches;
use std::f32;

struct RenderParams {
    flat: bool,
    iter: i32,
}

struct RenderBuffers(glium::VertexBuffer<RayVertex>, glium::IndexBuffer<u8>);
pub struct Renderer {
    program: glium::Program,
    background: glium::texture::SrgbTexture2d,

    buffers: RenderBuffers,

    params: RenderParams,
}

impl Renderer {
    pub fn new(display: &GlutinFacade, args: &ArgMatches) -> Self {
        let bg = {
            use std::io::Cursor;
            //FIXME: insert alternate bg images here

            let im = image::load(
                        Cursor::new(&include_bytes!("../resources/bg.jpg")[..]),
                        image::JPEG).unwrap().to_rgba();

            let imdim = im.dimensions();
            let im = glium::texture::RawImage2d::from_raw_rgba_reversed(
                        im.into_raw(), imdim);
            glium::texture::SrgbTexture2d::new(display, im).unwrap()
        };

        let prog = {
            //FIXME: add option for different shaders
            let shaders = &FLAT_SHADER;

            glium::Program::from_source(display,
                shaders.vert_shader,
                shaders.frag_shader,
                None).unwrap()
        };

        let bufs = {
            use glium::index::PrimitiveType::TrianglesList;
            RenderBuffers(glium::VertexBuffer::new(display, &VERTICES).unwrap(),
             glium::IndexBuffer::new(display, TrianglesList, &INDICES).unwrap())
        };

        Renderer { program: prog, background: bg, buffers: bufs,
            params: RenderParams::new(args) }
    }

    pub fn render(&self, display: &GlutinFacade, t: f32) {
        use glium::Surface;
        use time::precise_time_ns;

        let start = precise_time_ns();

        let mut target = display.draw();

        target.clear_color(0., 0., 0., 0.0);

        let (width, height) = target.get_dimensions();

        let src = [-3.0 * t.sin(),0.0,-3.0 * t.cos()];
        let facing_mat = {
            use cgmath::*;
            let dir = vec3(t.sin(), 0.0f32, t.cos());
            let up = vec3(0.,1.,0.0f32);

            // cgmath returns a tranposed look_at matrix for some reason
            Into::<[[f32;3];3]>::into(cgmath::Matrix3::look_at(dir, up)
                .transpose())
        };

        let (num_iter, time_step) = {
            let distance = 15.0f32; /* 2 * 15R_s, deflection is minimal by then */
            let num_iter = self.params.iter; /* arbitrary */

            (num_iter, distance / (num_iter as f32))
        };

        let uniforms = uniform! {
            height_ratio: (height as f32) / (width as f32),
            fov_ratio: (f32::consts::PI * 2. / 3. / 2.).tan(), // pi/2, 90 deg
            src: src,
            facing: facing_mat,
            tex: self.background
                .sampled().wrap_function(glium::uniforms::SamplerWrapFunction::Repeat),
            NUM_ITER: num_iter,
            TIME_STEP: time_step,
            FLAT: self.params.flat,
        };

        let params = glium::DrawParameters {
            blend: glium::Blend {
                color: glium::BlendingFunction::AlwaysReplace,
                alpha: glium::BlendingFunction::AlwaysReplace,
                constant_value: (0.0, 0.0, 0.0, 0.0),
            },
            .. Default::default()
        };

        target.draw(&self.buffers.0, &self.buffers.1, &self.program,
                    &uniforms, &params).unwrap();

        target.finish().unwrap();

        let end = precise_time_ns();
        println!("dt: {}ms", (end - start) as f32 / (1000000.0f32));
    }
}

impl RenderParams {
    fn new(args: &ArgMatches) -> Self {
        RenderParams {
            flat: args.is_present("flat"),
            iter: args.value_of("iter").unwrap().parse::<i32>().unwrap(),
        }
    }
}

#[derive(Copy, Clone)]
struct RayVertex {
    pos: (f32, f32),
}

implement_vertex!(RayVertex, pos);

const VERTICES: [RayVertex; 4] = [
    RayVertex { pos: (-1., -1.) },
    RayVertex { pos: (-1.,  1.) },
    RayVertex { pos: ( 1.,  1.) },
    RayVertex { pos: ( 1., -1.) },
];

const INDICES: [u8; 6] = [
    0, 1, 2,
    0, 2, 3
];

struct ShaderPair {
    vert_shader: &'static str,
    frag_shader: &'static str,
}

const FLAT_SHADER: ShaderPair = ShaderPair {
    vert_shader: r#"

#version 330

in vec2 pos;
out vec3 dir;
out vec2 pos_v;

uniform float height_ratio; // height / width
uniform float fov_ratio; // tan(fov / 2)

uniform mat3 facing;

void main() {
    float x = pos.x * fov_ratio;
    float y = pos.y * fov_ratio * height_ratio;
    dir = facing * vec3(x, y, 1.0);
    pos_v = pos;

    gl_Position = vec4(pos, 0.0, 1.0);
}

    "#,
    frag_shader: r#"

#version 330

#define M_PI 3.1415926535897932384626433832795

in vec3 dir;

out vec4 color;

uniform sampler2D bgtex;
uniform int NUM_ITER;
uniform float TIME_STEP;

/* we set constants to convenient values for now */
const float C = 1.0;
const float R_s = 1.0;
const float M = 0.5; /* must be R_s / 2 */
const float G = 1.0;

uniform bool FLAT;

uniform vec3 src;

float atan2(float y, float x) {
    return x == 0.0 ? sign(y) * M_PI / 2 : atan(y, x);
}

float yaw(vec3 v) {
    return atan2(v.x, v.z);
}

float yaw_coord(vec3 v) {
    return (yaw(v) + M_PI) / (2. * M_PI);
}

float pitch(vec3 v) {
    return asin(v.y);
}

float pitch_coord(vec3 v) {
    return (pitch(v) + M_PI / 2.) / M_PI;
}

vec4 bg_tex(vec3 dir) {
    float x = yaw_coord(dir);
    float y = pitch_coord(dir);

    vec2 tex_coords = vec2(x, y);

    float invert_x = x - 0.5;
    invert_x = invert_x - sign(invert_x) * 0.5;
    vec2 invert_coords = vec2(invert_x, y);

    vec2 dx1 = dFdx(tex_coords);
    vec2 dx2 = dFdx(invert_coords);

    vec2 dy1 = dFdy(tex_coords);
    vec2 dy2 = dFdy(invert_coords);

    vec2 dx = dot(dx1, dx1) < dot(dx2, dx2) ? dx1 : dx2;
    vec2 dy = dot(dy1, dy1) < dot(dy2, dy2) ? dy1 : dy2;

    /* force the LOD so that GLSL doesn't flip out on the discontinuity
       at the texture border */
    return textureGrad(bgtex, tex_coords, dx, dy);
}

const vec4 ZERO = vec4(0.0, 0.0, 0.0, 0.0);

void main() {
    vec3 ndir = normalize(dir);

    /* closest approach to BH */
    float dist = length(cross(ndir, src));

    /* test iteration function */
    int num_iter = int((1 / (1 + dist * 0.2)) * NUM_ITER);

    float alpha_rem = 1.0;
    vec4 ccolor = vec4(0.0, 0.0, 0.0, 0.0);
    vec3 cdir = ndir; /* current direction */

    vec3 pos = src;
    vec3 h = cross(pos, cdir);
    float h2 = dot(h, h); /* angular momentum */

    for(int i = 0; i < num_iter; i++) {
        vec3 npos = pos + cdir * TIME_STEP;
        if(!FLAT) {
            vec3 accel = -pos * 1.5 * h2 * pow(dot(pos, pos), -2.5);
            cdir = cdir + accel * TIME_STEP;
            if(0 == i % 100) {
                cdir = normalize(cdir);
                h = cross(pos, cdir);
                h2 = dot(h, h);
            }
            //cdir = cdir + accel * TIME_STEP;
        }

        /* check if its within a black hole */
        if(length(npos) <= R_s) {
            ccolor += alpha_rem * vec4(0.0, 0.0, 1.0, 1.0);
            alpha_rem *= 0.0;
        }

        pos = npos;
    }
    cdir = normalize(cdir);

    ccolor += alpha_rem * bg_tex(cdir);

    color = ccolor;
}

    "#,
};

/*const GR_SHADER: ShaderPair = ShaderPair {
    vert_shader: r#"

#version 140

in vec2 pos;
out vec2 pos_v;

void main() {
    pos_v = pos;

    gl_Position = vec4(pos, 0.0, 1.0);
}

    "#,
    frag_shader: r#"

#version 140

in vec2 pos_v;

out vec4 color;

void main() {
    float red = (pos_v.x + 1) / 2.;
    float green = (pos_v.y + 1) / 2.;
    color = vec4(red, green, 0.0, 1.0);
}

    "#,
};*/

<commit_msg>added new iteration scheme that lengthens time step for non-curving lines<commit_after>extern crate glium;
extern crate clap;
extern crate image;
extern crate cgmath;
extern crate time;

use glium::backend::glutin_backend::GlutinFacade;
use clap::ArgMatches;
use std::f32;

struct RenderParams {
    flat: bool,
    iter: i32,
}

struct RenderBuffers(glium::VertexBuffer<RayVertex>, glium::IndexBuffer<u8>);
pub struct Renderer {
    program: glium::Program,
    background: glium::texture::SrgbTexture2d,

    buffers: RenderBuffers,

    params: RenderParams,
}

impl Renderer {
    pub fn new(display: &GlutinFacade, args: &ArgMatches) -> Self {
        let bg = {
            use std::io::Cursor;
            //FIXME: insert alternate bg images here

            let im = image::load(
                        Cursor::new(&include_bytes!("../resources/bg.jpg")[..]),
                        image::JPEG).unwrap().to_rgba();

            let imdim = im.dimensions();
            let im = glium::texture::RawImage2d::from_raw_rgba_reversed(
                        im.into_raw(), imdim);
            glium::texture::SrgbTexture2d::new(display, im).unwrap()
        };

        let prog = {
            //FIXME: add option for different shaders
            let shaders = &FLAT_SHADER;

            glium::Program::from_source(display,
                shaders.vert_shader,
                shaders.frag_shader,
                None).unwrap()
        };

        let bufs = {
            use glium::index::PrimitiveType::TrianglesList;
            RenderBuffers(glium::VertexBuffer::new(display, &VERTICES).unwrap(),
             glium::IndexBuffer::new(display, TrianglesList, &INDICES).unwrap())
        };

        Renderer { program: prog, background: bg, buffers: bufs,
            params: RenderParams::new(args) }
    }

    pub fn render(&self, display: &GlutinFacade, t: f32) {
        use glium::Surface;
        use time::precise_time_ns;

        let start = precise_time_ns();

        let mut target = display.draw();

        target.clear_color(0., 0., 0., 0.0);

        let (width, height) = target.get_dimensions();

        let src = [-10.0 * t.sin(),0.0,-10.0 * t.cos()];
        let facing_mat = {
            use cgmath::*;
            let dir = vec3(t.sin(), 0.0f32, t.cos());
            let up = vec3(0.,1.,0.0f32);

            // cgmath returns a tranposed look_at matrix for some reason
            Into::<[[f32;3];3]>::into(cgmath::Matrix3::look_at(dir, up)
                .transpose())
        };

        let (num_iter, time_step) = {
            let distance = 15.0f32; /* 2 * 15R_s, deflection is minimal by then */
            let num_iter = self.params.iter; /* arbitrary */

            (num_iter, distance / (num_iter as f32))
        };

        let uniforms = uniform! {
            height_ratio: (height as f32) / (width as f32),
            fov_ratio: (f32::consts::PI * 2. / 3. / 2.).tan(), // pi/2, 90 deg
            src: src,
            facing: facing_mat,
            tex: self.background
                .sampled().wrap_function(glium::uniforms::SamplerWrapFunction::Repeat),
            NUM_ITER: num_iter,
            TIME_STEP: time_step,
            FLAT: self.params.flat,
        };

        let params = glium::DrawParameters {
            blend: glium::Blend {
                color: glium::BlendingFunction::AlwaysReplace,
                alpha: glium::BlendingFunction::AlwaysReplace,
                constant_value: (0.0, 0.0, 0.0, 0.0),
            },
            .. Default::default()
        };

        target.draw(&self.buffers.0, &self.buffers.1, &self.program,
                    &uniforms, &params).unwrap();

        target.finish().unwrap();

        let end = precise_time_ns();
        println!("dt: {}ms", (end - start) as f32 / (1000000.0f32));
    }
}

impl RenderParams {
    fn new(args: &ArgMatches) -> Self {
        RenderParams {
            flat: args.is_present("flat"),
            iter: args.value_of("iter").unwrap().parse::<i32>().unwrap(),
        }
    }
}

#[derive(Copy, Clone)]
struct RayVertex {
    pos: (f32, f32),
}

implement_vertex!(RayVertex, pos);

const VERTICES: [RayVertex; 4] = [
    RayVertex { pos: (-1., -1.) },
    RayVertex { pos: (-1.,  1.) },
    RayVertex { pos: ( 1.,  1.) },
    RayVertex { pos: ( 1., -1.) },
];

const INDICES: [u8; 6] = [
    0, 1, 2,
    0, 2, 3
];

struct ShaderPair {
    vert_shader: &'static str,
    frag_shader: &'static str,
}

const FLAT_SHADER: ShaderPair = ShaderPair {
    vert_shader: r#"

#version 330

in vec2 pos;
out vec3 dir;
out vec2 pos_v;

uniform float height_ratio; // height / width
uniform float fov_ratio; // tan(fov / 2)

uniform mat3 facing;

void main() {
    float x = pos.x * fov_ratio;
    float y = pos.y * fov_ratio * height_ratio;
    dir = facing * vec3(x, y, 1.0);
    pos_v = pos;

    gl_Position = vec4(pos, 0.0, 1.0);
}

    "#,
    frag_shader: r#"

#version 330

#define M_PI 3.1415926535897932384626433832795

in vec3 dir;

out vec4 color;

uniform sampler2D bgtex;
uniform int NUM_ITER;
uniform float TIME_STEP;

/* we set constants to convenient values for now */
const float C = 1.0;
const float R_s = 1.0;
const float M = 0.5; /* must be R_s / 2 */
const float G = 1.0;

uniform bool FLAT;

uniform vec3 src;

float atan2(float y, float x) {
    return x == 0.0 ? sign(y) * M_PI / 2 : atan(y, x);
}

float yaw(vec3 v) {
    return atan2(v.x, v.z);
}

float yaw_coord(vec3 v) {
    return (yaw(v) + M_PI) / (2. * M_PI);
}

float pitch(vec3 v) {
    return asin(v.y);
}

float pitch_coord(vec3 v) {
    return (pitch(v) + M_PI / 2.) / M_PI;
}

vec4 bg_tex(vec3 dir) {
    float x = yaw_coord(dir);
    float y = pitch_coord(dir);

    vec2 tex_coords = vec2(x, y);

    float invert_x = x - 0.5;
    invert_x = invert_x - sign(invert_x) * 0.5;
    vec2 invert_coords = vec2(invert_x, y);

    vec2 dx1 = dFdx(tex_coords);
    vec2 dx2 = dFdx(invert_coords);

    vec2 dy1 = dFdy(tex_coords);
    vec2 dy2 = dFdy(invert_coords);

    vec2 dx = dot(dx1, dx1) < dot(dx2, dx2) ? dx1 : dx2;
    vec2 dy = dot(dy1, dy1) < dot(dy2, dy2) ? dy1 : dy2;

    /* force the LOD so that GLSL doesn't flip out on the discontinuity
       at the texture border */
    return textureGrad(bgtex, tex_coords, dx, dy);
}

const vec4 ZERO = vec4(0.0, 0.0, 0.0, 0.0);
const float TDIST = 20.0;
const float BTHRESHOLD = 0.9;
const float TTHRESHOLD = 0.95;
const float M_RAT = 8;

void main() {
    float cdist = 0.0;
    float ratio = 1.0;
    vec3 ndir = normalize(dir);

    /* closest approach to BH */
    float dist = length(cross(ndir, src));

    /* test iteration function */
    int num_iter = int((1 / (1 + dist * 0.2)) * NUM_ITER);

    float alpha_rem = 1.0;
    vec4 ccolor = vec4(0.0, 0.0, 0.0, 0.0);
    vec3 cdir = ndir; /* current direction */

    vec3 pos = src;
    vec3 h = cross(pos, cdir);
    float h2 = dot(h, h); /* angular momentum */

    while(cdist < TDIST) {
        vec3 npos = pos + cdir * TIME_STEP * ratio;
        cdist += TIME_STEP * ratio;
        if(!FLAT) {
            vec3 accel = -pos * 1.5 * h2 * pow(dot(pos, pos), -2.5);
            vec3 ncdir = cdir + accel * TIME_STEP * ratio;
            ncdir = normalize(ncdir);
            h = cross(pos, ncdir);
            h2 = dot(h, h);
            if(dot(ncdir, cdir) < BTHRESHOLD && ratio > (1/M_RAT)) {
                ratio *= 0.5;
            } else if(dot(ncdir, cdir) > TTHRESHOLD && ratio < M_RAT) {
                ratio *= 2;
            }
            cdir = ncdir;
        }

        /* check if its within a black hole */
        if(length(npos) <= R_s) {
            ccolor += alpha_rem * vec4(0.0, 0.0, 0.0, 1.0);
            alpha_rem *= 0.0;
        }

        pos = npos;
    }
/*
    for(int i = 0; i < num_iter; i++) {
        vec3 npos = pos + cdir * TIME_STEP;
        if(!FLAT) {
            vec3 accel = -pos * 1.5 * h2 * pow(dot(pos, pos), -2.5);
            cdir = cdir + accel * TIME_STEP;
            cdir = normalize(cdir);
            h = cross(pos, cdir);
            h2 = dot(h, h);
            //cdir = cdir + accel * TIME_STEP;
        }

        /* check if its within a black hole */
  /*      if(length(npos) <= R_s) {
            ccolor += alpha_rem * vec4(0.0, 0.0, 1.0, 1.0);
            alpha_rem *= 0.0;
        }

        pos = npos;
    }*/
    cdir = normalize(cdir);

    ccolor += alpha_rem * bg_tex(cdir);

    color = ccolor;
}

    "#,
};

/*const GR_SHADER: ShaderPair = ShaderPair {
    vert_shader: r#"

#version 140

in vec2 pos;
out vec2 pos_v;

void main() {
    pos_v = pos;

    gl_Position = vec4(pos, 0.0, 1.0);
}

    "#,
    frag_shader: r#"

#version 140

in vec2 pos_v;

out vec4 color;

void main() {
    float red = (pos_v.x + 1) / 2.;
    float green = (pos_v.y + 1) / 2.;
    color = vec4(red, green, 0.0, 1.0);
}

    "#,
};*/

<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::env;
use std::path::Path;

use crate::stacked_borrows::Tag;
use crate::*;
use rustc::ty::layout::Size;
use rustc_mir::interpret::{Memory, Pointer};

#[derive(Default)]
pub struct EnvVars {
    /// Stores pointers to the environment variables. These variables must be stored as
    /// null-terminated C strings with the `"{name}={value}"` format.
    map: HashMap<Vec<u8>, Pointer<Tag>>,
}

impl EnvVars {
    pub(crate) fn init<'mir, 'tcx>(
        ecx: &mut InterpCx<'mir, 'tcx, Evaluator<'tcx>>,
        mut excluded_env_vars: Vec<String>,
    ) {
        // Exclude `TERM` var to avoid terminfo trying to open the termcap file.
        excluded_env_vars.push("TERM".to_owned());

        if ecx.machine.communicate {
            for (name, value) in env::vars() {
                if !excluded_env_vars.contains(&name) {
                    let var_ptr =
                        alloc_env_var(name.as_bytes(), value.as_bytes(), ecx.memory_mut());
                    ecx.machine.env_vars.map.insert(name.into_bytes(), var_ptr);
                }
            }
        }
    }
}

fn alloc_env_var<'mir, 'tcx>(
    name: &[u8],
    value: &[u8],
    memory: &mut Memory<'mir, 'tcx, Evaluator<'tcx>>,
) -> Pointer<Tag> {
    let mut bytes = name.to_vec();
    bytes.push(b'=');
    bytes.extend_from_slice(value);
    bytes.push(0);
    memory.allocate_static_bytes(bytes.as_slice(), MiriMemoryKind::Env.into())
}

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn getenv(&mut self, name_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, Scalar<Tag>> {
        let this = self.eval_context_mut();

        let name_ptr = this.read_scalar(name_op)?.not_undef()?;
        let name = this.memory().read_c_str(name_ptr)?;
        Ok(match this.machine.env_vars.map.get(name) {
            // The offset is used to strip the "{name}=" part of the string.
            Some(var_ptr) => {
                Scalar::Ptr(var_ptr.offset(Size::from_bytes(name.len() as u64 + 1), this)?)
            }
            None => Scalar::ptr_null(&*this.tcx),
        })
    }

    fn setenv(
        &mut self,
        name_op: OpTy<'tcx, Tag>,
        value_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        let name_ptr = this.read_scalar(name_op)?.not_undef()?;
        let value_ptr = this.read_scalar(value_op)?.not_undef()?;
        let value = this.memory().read_c_str(value_ptr)?;
        let mut new = None;
        if !this.is_null(name_ptr)? {
            let name = this.memory().read_c_str(name_ptr)?;
            if !name.is_empty() && !name.contains(&b'=') {
                new = Some((name.to_owned(), value.to_owned()));
            }
        }
        if let Some((name, value)) = new {
            let var_ptr = alloc_env_var(&name, &value, this.memory_mut());
            if let Some(var) = this.machine.env_vars.map.insert(name.to_owned(), var_ptr) {
                this.memory_mut()
                    .deallocate(var, None, MiriMemoryKind::Env.into())?;
            }
            Ok(0)
        } else {
            Ok(-1)
        }
    }

    fn unsetenv(&mut self, name_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        let name_ptr = this.read_scalar(name_op)?.not_undef()?;
        let mut success = None;
        if !this.is_null(name_ptr)? {
            let name = this.memory().read_c_str(name_ptr)?.to_owned();
            if !name.is_empty() && !name.contains(&b'=') {
                success = Some(this.machine.env_vars.map.remove(&name));
            }
        }
        if let Some(old) = success {
            if let Some(var) = old {
                this.memory_mut()
                    .deallocate(var, None, MiriMemoryKind::Env.into())?;
            }
            Ok(0)
        } else {
            Ok(-1)
        }
    }

    fn getcwd(
        &mut self,
        buf_op: OpTy<'tcx, Tag>,
        size_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, Scalar<Tag>> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`getcwd` not available when isolation is enabled")
        }

        let tcx = &{ this.tcx.tcx };

        let buf = this.force_ptr(this.read_scalar(buf_op)?.not_undef()?)?;
        let size = this.read_scalar(size_op)?.to_usize(&*this.tcx)?;
        // If we cannot get the current directory, we return null
        match env::current_dir() {
            Ok(cwd) => {
                // It is not clear what happens with non-utf8 paths here
                let mut bytes = cwd.display().to_string().into_bytes();
                // If `size` is smaller or equal than the `bytes.len()`, writing `bytes` plus the
                // required null terminator to memory using the `buf` pointer would cause an
                // overflow, the desired behavior in this case is to return null.
                if (bytes.len() as u64) < size {
                    // We add a `/0` terminator
                    bytes.push(0);
                    // This is ok because the buffer was strictly larger than `bytes`, so after
                    // adding the null terminator, the buffer size is larger or equal to
                    // `bytes.len()`, meaning that `bytes` actually fit inside tbe buffer.
                    this.memory_mut()
                        .get_mut(buf.alloc_id)?
                        .write_bytes(tcx, buf, &bytes)?;
                    return Ok(Scalar::Ptr(buf));
                }
                let erange = this.eval_libc("ERANGE")?;
                this.set_last_error(erange)?;
            }
            Err(e) => this.consume_io_error(e)?,
        }
        Ok(Scalar::ptr_null(&*this.tcx))
    }

    fn chdir(&mut self, path_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`chdir` not available when isolation is enabled")
        }

        let path_bytes = this
            .memory()
            .read_c_str(this.read_scalar(path_op)?.not_undef()?)?;

        let path = Path::new(
            std::str::from_utf8(path_bytes)
                .map_err(|_| err_unsup_format!("{:?} is not a valid utf-8 string", path_bytes))?,
        );

        match env::set_current_dir(path) {
            Ok(()) => Ok(0),
            Err(e) => {
                this.consume_io_error(e)?;
                Ok(-1)
            }
        }
    }
}
<commit_msg>Use existing tcx instead<commit_after>use std::collections::HashMap;
use std::env;
use std::path::Path;

use crate::stacked_borrows::Tag;
use crate::*;
use rustc::ty::layout::Size;
use rustc_mir::interpret::{Memory, Pointer};

#[derive(Default)]
pub struct EnvVars {
    /// Stores pointers to the environment variables. These variables must be stored as
    /// null-terminated C strings with the `"{name}={value}"` format.
    map: HashMap<Vec<u8>, Pointer<Tag>>,
}

impl EnvVars {
    pub(crate) fn init<'mir, 'tcx>(
        ecx: &mut InterpCx<'mir, 'tcx, Evaluator<'tcx>>,
        mut excluded_env_vars: Vec<String>,
    ) {
        // Exclude `TERM` var to avoid terminfo trying to open the termcap file.
        excluded_env_vars.push("TERM".to_owned());

        if ecx.machine.communicate {
            for (name, value) in env::vars() {
                if !excluded_env_vars.contains(&name) {
                    let var_ptr =
                        alloc_env_var(name.as_bytes(), value.as_bytes(), ecx.memory_mut());
                    ecx.machine.env_vars.map.insert(name.into_bytes(), var_ptr);
                }
            }
        }
    }
}

fn alloc_env_var<'mir, 'tcx>(
    name: &[u8],
    value: &[u8],
    memory: &mut Memory<'mir, 'tcx, Evaluator<'tcx>>,
) -> Pointer<Tag> {
    let mut bytes = name.to_vec();
    bytes.push(b'=');
    bytes.extend_from_slice(value);
    bytes.push(0);
    memory.allocate_static_bytes(bytes.as_slice(), MiriMemoryKind::Env.into())
}

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn getenv(&mut self, name_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, Scalar<Tag>> {
        let this = self.eval_context_mut();

        let name_ptr = this.read_scalar(name_op)?.not_undef()?;
        let name = this.memory().read_c_str(name_ptr)?;
        Ok(match this.machine.env_vars.map.get(name) {
            // The offset is used to strip the "{name}=" part of the string.
            Some(var_ptr) => {
                Scalar::Ptr(var_ptr.offset(Size::from_bytes(name.len() as u64 + 1), this)?)
            }
            None => Scalar::ptr_null(&*this.tcx),
        })
    }

    fn setenv(
        &mut self,
        name_op: OpTy<'tcx, Tag>,
        value_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        let name_ptr = this.read_scalar(name_op)?.not_undef()?;
        let value_ptr = this.read_scalar(value_op)?.not_undef()?;
        let value = this.memory().read_c_str(value_ptr)?;
        let mut new = None;
        if !this.is_null(name_ptr)? {
            let name = this.memory().read_c_str(name_ptr)?;
            if !name.is_empty() && !name.contains(&b'=') {
                new = Some((name.to_owned(), value.to_owned()));
            }
        }
        if let Some((name, value)) = new {
            let var_ptr = alloc_env_var(&name, &value, this.memory_mut());
            if let Some(var) = this.machine.env_vars.map.insert(name.to_owned(), var_ptr) {
                this.memory_mut()
                    .deallocate(var, None, MiriMemoryKind::Env.into())?;
            }
            Ok(0)
        } else {
            Ok(-1)
        }
    }

    fn unsetenv(&mut self, name_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        let name_ptr = this.read_scalar(name_op)?.not_undef()?;
        let mut success = None;
        if !this.is_null(name_ptr)? {
            let name = this.memory().read_c_str(name_ptr)?.to_owned();
            if !name.is_empty() && !name.contains(&b'=') {
                success = Some(this.machine.env_vars.map.remove(&name));
            }
        }
        if let Some(old) = success {
            if let Some(var) = old {
                this.memory_mut()
                    .deallocate(var, None, MiriMemoryKind::Env.into())?;
            }
            Ok(0)
        } else {
            Ok(-1)
        }
    }

    fn getcwd(
        &mut self,
        buf_op: OpTy<'tcx, Tag>,
        size_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, Scalar<Tag>> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`getcwd` not available when isolation is enabled")
        }

        let tcx = &{ this.tcx.tcx };

        let buf = this.force_ptr(this.read_scalar(buf_op)?.not_undef()?)?;
        let size = this.read_scalar(size_op)?.to_usize(&*tcx)?;
        // If we cannot get the current directory, we return null
        match env::current_dir() {
            Ok(cwd) => {
                // It is not clear what happens with non-utf8 paths here
                let mut bytes = cwd.display().to_string().into_bytes();
                // If `size` is smaller or equal than the `bytes.len()`, writing `bytes` plus the
                // required null terminator to memory using the `buf` pointer would cause an
                // overflow, the desired behavior in this case is to return null.
                if (bytes.len() as u64) < size {
                    // We add a `/0` terminator
                    bytes.push(0);
                    // This is ok because the buffer was strictly larger than `bytes`, so after
                    // adding the null terminator, the buffer size is larger or equal to
                    // `bytes.len()`, meaning that `bytes` actually fit inside tbe buffer.
                    this.memory_mut()
                        .get_mut(buf.alloc_id)?
                        .write_bytes(tcx, buf, &bytes)?;
                    return Ok(Scalar::Ptr(buf));
                }
                let erange = this.eval_libc("ERANGE")?;
                this.set_last_error(erange)?;
            }
            Err(e) => this.consume_io_error(e)?,
        }
        Ok(Scalar::ptr_null(&*tcx))
    }

    fn chdir(&mut self, path_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`chdir` not available when isolation is enabled")
        }

        let path_bytes = this
            .memory()
            .read_c_str(this.read_scalar(path_op)?.not_undef()?)?;

        let path = Path::new(
            std::str::from_utf8(path_bytes)
                .map_err(|_| err_unsup_format!("{:?} is not a valid utf-8 string", path_bytes))?,
        );

        match env::set_current_dir(path) {
            Ok(()) => Ok(0),
            Err(e) => {
                this.consume_io_error(e)?;
                Ok(-1)
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::io;

use std::fmt::Write;

use chrono::prelude::{Local, DateTime};

use pombase_rc_string::RcString;

use crate::web::config::*;
use crate::data_types::*;

pub const GO_ASPECT_NAMES: [&str; 3] =
    ["cellular_component", "biological_process", "molecular_function"];

fn eco_evidence_from_annotation(mapping: &GoEcoMapping,
                                annotation: &OntAnnotationDetail)
                                -> String
{
    if let Some(ref go_evidence) = annotation.evidence {
        if let Some(ref reference) = annotation.reference {
            mapping.lookup_with_go_ref(go_evidence, reference)
                .unwrap_or_else(|| mapping.lookup_default(go_evidence)
                                .unwrap_or_else(|| panic!("failed to find {} for {} in GO mapping",
                                                          go_evidence, reference)))
        } else {
            mapping.lookup_default(go_evidence)
                .unwrap_or_else(|| panic!("failed to find {} in GO mapping",
                                          go_evidence))
        }
    } else {
        String::from("")
    }
}

fn get_gpad_relation_of(term_details: &TermDetails,
                        annotation_detail: &OntAnnotationDetail) -> String {
    match term_details.cv_name.as_str() {
        "molecular_function" => {
            if annotation_detail.qualifiers.iter()
                .find(|q| q.as_str() == "contributes_to").is_some()
            {
                String::from("RO:0002326")
            } else {
                String::from("RO:0002327")
            }
        },
        "biological_process" => String::from("RO:0002331"),
        "cellular_component" => {
            if term_details.interesting_parent_details.iter()
                .any(|p| p.termid == RcString::from("GO:0032991"))
            {
                String::from("BFO:0000050")
            } else {
                String::from("RO:0001025")
            }
        },
        _ => panic!("unknown cv_name in GPAD export: {}", term_details.cv_name)
    }
}

fn make_gpad_extension_string(config: &Config, extension: &Vec<ExtPart>) -> String {
    let rel_mapping = &config.file_exports.gpad_gpi.extension_relation_mappings;
    let get_rel_termid = |ext_part: &ExtPart| {
        if let Some(map_termid) = rel_mapping.get(ext_part.rel_type_name.as_ref()) {
            map_termid.clone().unwrap()
        } else {
            ext_part.rel_type_id.clone().unwrap()
        }
    };

    let get_range = |ext_part: &ExtPart| {
        let mut range_copy = ext_part.ext_range.clone();

        if let ExtRange::Gene(ref mut gene_uniquename) = range_copy {
            if !gene_uniquename.contains(":") {
                let new_uniquename =
                    RcString::from(&format!("{}:{}", config.database_name,
                                            gene_uniquename));
                *gene_uniquename = new_uniquename;
            }
        }
        range_copy
    };

    extension.iter()
        .filter(|ext_part| {
            if let Some(map_termid) = rel_mapping.get(ext_part.rel_type_name.as_ref()) {
                map_termid.is_some()
            } else {
                ext_part.rel_type_id.is_some()
            }
        })
        .map(|ext_part| format!("{}({})", get_rel_termid(ext_part),
                                get_range(ext_part)))
        .collect::<Vec<_>>().join(",")
}

pub fn write_gene_product_annotation(gpad_writer: &mut dyn io::Write,
                                     go_eco_mappping: &GoEcoMapping, config: &Config,
                                     api_maps: &APIMaps, gene_details: &GeneDetails)
                                     -> Result<(), io::Error>
{
    let database_name = &config.database_name;
    let db_object_id = format!("{}:{}", database_name, gene_details.uniquename);
    let local: DateTime<Local> = Local::now();
    let local_iso_date = local.format("%F");
    let assigned_by = &config.database_name;

    if gene_details.feature_type == "ncRNA gene" {
        let mut found = false;
        for aspect in GO_ASPECT_NAMES.iter() {
            let term_annotations = gene_details.cv_annotations.get(&RcString::from(aspect));
            if term_annotations.is_some() {
                found = true;
                break;
            }
        }
        if !found {
            // special case for ncRNA genes: only write ND lines for an aspect
            // if there are some annotations for the other aspects
            return Ok(());
        }
    }

    for aspect in GO_ASPECT_NAMES.iter() {
        let term_annotations = gene_details.cv_annotations.get(&RcString::from(aspect));
        if term_annotations.is_none() {
            let go_aspect_termid =
                config.file_exports.gpad_gpi.go_aspect_terms.get(*aspect).unwrap();
            let nd_ref = &config.file_exports.nd_reference;
            let line = format!("{}\t\tsome_rel\t{}\t{}\tECO:0000307\t\t\t{}\t{}\t\t\n",
                               db_object_id,
                               &go_aspect_termid,
                               nd_ref,
                               local_iso_date, assigned_by);
            gpad_writer.write_all(line.as_bytes())?;
        }
    }

    for (cv_name, term_annotations) in &gene_details.cv_annotations {
        if !GO_ASPECT_NAMES.contains(&cv_name.as_ref()) {
            continue;
        }

        for term_annotation in term_annotations {
            let term = api_maps.terms.get(&term_annotation.term)
                .expect(&format!("failed to find term summary for {}",
                                 term_annotation.term));

            for annotation_id in &term_annotation.annotations {
                let annotation_detail = api_maps.annotation_details
                    .get(&annotation_id)
                    .expect(&format!("can't find annotation {}", annotation_id));
                let not = if term_annotation.is_not {
                    "NOT"
                } else {
                    ""
                };

                let relation = get_gpad_relation_of(term, annotation_detail);

                let ontology_class_id = &term_annotation.term;
                let reference_uniquename =
                    annotation_detail.reference.clone()
                    .unwrap_or_else(|| RcString::from(""));
                let evidence_type =
                    eco_evidence_from_annotation(go_eco_mappping, annotation_detail);

                let with_iter = annotation_detail.withs.iter();
                let from_iter = annotation_detail.froms.iter();
                let mut with_or_from_parts =
                    with_iter.chain(from_iter).map(|s| {
                        let s: RcString = s.clone().into();
                        if s.contains(":") {
                            s
                        } else {
                            RcString::from(&format!("{}:{}", assigned_by, s))
                        }
                    })
                    .collect::<Vec<RcString>>();
                with_or_from_parts.sort_unstable();
                let with_or_from = with_or_from_parts.join(",");

                if let Some(ref date) = annotation_detail.date {
                    let annotation_extensions =
                        make_gpad_extension_string(config, &annotation_detail.extension);
                    let line = format!("{}\t{}\t{}\t{}\t{}\t{}\t{}\t\t{}\t{}\t{}\t\n",
                                       db_object_id,
                                       not, relation, ontology_class_id,
                                       reference_uniquename, evidence_type,
                                       with_or_from, date, assigned_by,
                                       annotation_extensions);
                    gpad_writer.write_all(line.as_bytes())?;
                } else {
                    println!("WARNING while writing GAF file: \
                              date missing from annotation {}", annotation_id);
                }
            }
        }
    }

    Ok(())
}

fn make_gaf_extension_string(config: &Config, extension: &Vec<ExtPart>) -> String {
    let mut ret_string = String::new();

    for (idx, ext_part) in extension.iter().enumerate() {
        ret_string += &ext_part.rel_type_name;
        ret_string.push('(');

        if let ExtRange::Gene(ref gene_uniquename) = ext_part.ext_range {
            if gene_uniquename.contains(":") {
                ret_string += gene_uniquename;
            } else {
                ret_string += &config.database_name;
                ret_string.push(':');
                ret_string += &gene_uniquename;
            }
        } else {
            write!(ret_string, "{}", &ext_part.ext_range).unwrap();
        }
        ret_string.push(')');

        if idx < extension.len() - 1 {
            ret_string.push(',');
        }
    }

    ret_string
}

pub fn write_go_annotation_format(writer: &mut dyn io::Write, config: &Config,
                                  api_maps: &APIMaps, gene_details: &GeneDetails,
                                  cv_name: &str)
                                  -> Result<(), io::Error>
{
    if !GO_ASPECT_NAMES.contains(&cv_name) {
        return Err(io::Error::new(io::ErrorKind::InvalidInput,
                                  format!("unknown CV: {}", cv_name)));
    }

    let database_name = &config.database_name;
    let db_object_id = &gene_details.uniquename;
    let db_object_symbol =
        if let Some(ref gene_name) = gene_details.name {
            gene_name.as_str()
        } else {
            &db_object_id
        };
    let db_object_synonyms =
        gene_details.synonyms.iter().filter(|synonym| {
            synonym.synonym_type == "exact"
        })
        .map(|synonym| synonym.name.to_string())
        .collect::<Vec<String>>()
        .join("|");

    let single_letter_aspect =
        if cv_name.starts_with("b") {
            "P"
        } else {
            if cv_name.starts_with("c") {
                "C"
            } else {
                "F"
            }
        };

    if let Some(term_annotations) = gene_details.cv_annotations.get(cv_name) {
        for term_annotation in term_annotations {
            for annotation_id in &term_annotation.annotations {
                let annotation_detail = api_maps.annotation_details
                    .get(&annotation_id)
                    .expect(&format!("can't find annotation {}", annotation_id));

                let mut qualifier_parts = vec![];
                if term_annotation.is_not {
                    qualifier_parts.push("NOT");
                };

                qualifier_parts.extend(annotation_detail.qualifiers.iter()
                                       .map(|s| s.as_str()));

                let go_id = &term_annotation.term;
                let reference_uniquename =
                    annotation_detail.reference.clone()
                    .unwrap_or_else(|| RcString::from(""));
                let assigned_by =
                    if let Some(ref assigned_by) = annotation_detail.assigned_by {
                        assigned_by.as_str()
                    } else {
                        database_name.as_str()
                    };

                let with_iter = annotation_detail.withs.iter();
                let from_iter = annotation_detail.froms.iter();
                let mut with_or_from_parts =
                    with_iter.chain(from_iter).map(|s| {
                        let s: RcString = s.clone().into();
                        if s.contains(":") {
                            s
                        } else {
                            RcString::from(&format!("{}:{}", assigned_by, s))
                        }
                    })
                    .collect::<Vec<RcString>>();
                with_or_from_parts.sort_unstable();
                let with_or_from = with_or_from_parts.join(",");

                let evidence_code =
                    annotation_detail.evidence
                    .as_ref()
                    .map(|s| s.as_str())
                    .unwrap_or_else(|| "");

                let db_object_name =
                    if let Some(ref product) = gene_details.product {
                        product.as_str()
                    } else {
                        ""
                    };

                let transcript_type_str =
                    gene_details.transcripts[0].transcript_type.as_str();
                let db_object_type =
                    if transcript_type_str == "mRNA" {
                        "protein"
                    } else {
                        transcript_type_str
                    };
                let date =
                    if let Some(ref raw_date) = annotation_detail.date {
                        raw_date.replace("-", "")
                    } else {
                        println!("WARNING while writing GPAD file: \
                                  date missing from annotation {}", annotation_id);
                        continue;
                    };

                let annotation_extensions =
                    make_gaf_extension_string(config, &annotation_detail.extension);

                let gene_product_form_id =
                    if let Some(ref gene_product_form_id) = annotation_detail.gene_product_form_id
                {
                    gene_product_form_id.as_str()
                } else {
                    ""
                };

                let line = format!("{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\ttaxon:{}\t{}\t{}\t{}\t{}\n",
                                   database_name,
                                   db_object_id,
                                   db_object_symbol,
                                   qualifier_parts.join("|"),
                                   go_id,
                                   reference_uniquename,
                                   evidence_code,
                                   with_or_from,
                                   single_letter_aspect,
                                   db_object_name,
                                   db_object_synonyms,
                                   db_object_type,
                                   gene_details.taxonid,
                                   date,
                                   assigned_by,
                                   annotation_extensions,
                                   gene_product_form_id);
                writer.write_all(line.as_bytes())?;
            }
        }
    }

    Ok(())
}

<commit_msg>Don't write dubious or pseudogene lines to GPAD<commit_after>use std::io;

use std::fmt::Write;

use chrono::prelude::{Local, DateTime};

use pombase_rc_string::RcString;

use crate::web::config::*;
use crate::data_types::*;

pub const GO_ASPECT_NAMES: [&str; 3] =
    ["cellular_component", "biological_process", "molecular_function"];

fn eco_evidence_from_annotation(mapping: &GoEcoMapping,
                                annotation: &OntAnnotationDetail)
                                -> String
{
    if let Some(ref go_evidence) = annotation.evidence {
        if let Some(ref reference) = annotation.reference {
            mapping.lookup_with_go_ref(go_evidence, reference)
                .unwrap_or_else(|| mapping.lookup_default(go_evidence)
                                .unwrap_or_else(|| panic!("failed to find {} for {} in GO mapping",
                                                          go_evidence, reference)))
        } else {
            mapping.lookup_default(go_evidence)
                .unwrap_or_else(|| panic!("failed to find {} in GO mapping",
                                          go_evidence))
        }
    } else {
        String::from("")
    }
}

fn get_gpad_relation_of(term_details: &TermDetails,
                        annotation_detail: &OntAnnotationDetail) -> String {
    match term_details.cv_name.as_str() {
        "molecular_function" => {
            if annotation_detail.qualifiers.iter()
                .find(|q| q.as_str() == "contributes_to").is_some()
            {
                String::from("RO:0002326")
            } else {
                String::from("RO:0002327")
            }
        },
        "biological_process" => String::from("RO:0002331"),
        "cellular_component" => {
            if term_details.interesting_parent_details.iter()
                .any(|p| p.termid == RcString::from("GO:0032991"))
            {
                String::from("BFO:0000050")
            } else {
                String::from("RO:0001025")
            }
        },
        _ => panic!("unknown cv_name in GPAD export: {}", term_details.cv_name)
    }
}

fn make_gpad_extension_string(config: &Config, extension: &Vec<ExtPart>) -> String {
    let rel_mapping = &config.file_exports.gpad_gpi.extension_relation_mappings;
    let get_rel_termid = |ext_part: &ExtPart| {
        if let Some(map_termid) = rel_mapping.get(ext_part.rel_type_name.as_ref()) {
            map_termid.clone().unwrap()
        } else {
            ext_part.rel_type_id.clone().unwrap()
        }
    };

    let get_range = |ext_part: &ExtPart| {
        let mut range_copy = ext_part.ext_range.clone();

        if let ExtRange::Gene(ref mut gene_uniquename) = range_copy {
            if !gene_uniquename.contains(":") {
                let new_uniquename =
                    RcString::from(&format!("{}:{}", config.database_name,
                                            gene_uniquename));
                *gene_uniquename = new_uniquename;
            }
        }
        range_copy
    };

    extension.iter()
        .filter(|ext_part| {
            if let Some(map_termid) = rel_mapping.get(ext_part.rel_type_name.as_ref()) {
                map_termid.is_some()
            } else {
                ext_part.rel_type_id.is_some()
            }
        })
        .map(|ext_part| format!("{}({})", get_rel_termid(ext_part),
                                get_range(ext_part)))
        .collect::<Vec<_>>().join(",")
}

pub fn write_gene_product_annotation(gpad_writer: &mut dyn io::Write,
                                     go_eco_mappping: &GoEcoMapping, config: &Config,
                                     api_maps: &APIMaps, gene_details: &GeneDetails)
                                     -> Result<(), io::Error>
{
    let database_name = &config.database_name;
    let db_object_id = format!("{}:{}", database_name, gene_details.uniquename);
    let local: DateTime<Local> = Local::now();
    let local_iso_date = local.format("%F");
    let assigned_by = &config.database_name;

    if let Some(ref product) = gene_details.product {
        if product == "dubious" {
            return Ok(());
        }
    }

    if gene_details.feature_type == "pseudogene" {
        return Ok(());
    }

    if gene_details.feature_type == "ncRNA gene" {
        let mut found = false;
        for aspect in GO_ASPECT_NAMES.iter() {
            let term_annotations = gene_details.cv_annotations.get(&RcString::from(aspect));
            if term_annotations.is_some() {
                found = true;
                break;
            }
        }
        if !found {
            // special case for ncRNA genes: only write ND lines for an aspect
            // if there are some annotations for the other aspects
            return Ok(());
        }
    }

    for aspect in GO_ASPECT_NAMES.iter() {
        let term_annotations = gene_details.cv_annotations.get(&RcString::from(aspect));
        if term_annotations.is_none() {
            let go_aspect_termid =
                config.file_exports.gpad_gpi.go_aspect_terms.get(*aspect).unwrap();
            let nd_ref = &config.file_exports.nd_reference;
            let line = format!("{}\t\tsome_rel\t{}\t{}\tECO:0000307\t\t\t{}\t{}\t\t\n",
                               db_object_id,
                               &go_aspect_termid,
                               nd_ref,
                               local_iso_date, assigned_by);
            gpad_writer.write_all(line.as_bytes())?;
        }
    }

    for (cv_name, term_annotations) in &gene_details.cv_annotations {
        if !GO_ASPECT_NAMES.contains(&cv_name.as_ref()) {
            continue;
        }

        for term_annotation in term_annotations {
            let term = api_maps.terms.get(&term_annotation.term)
                .expect(&format!("failed to find term summary for {}",
                                 term_annotation.term));

            for annotation_id in &term_annotation.annotations {
                let annotation_detail = api_maps.annotation_details
                    .get(&annotation_id)
                    .expect(&format!("can't find annotation {}", annotation_id));
                let not = if term_annotation.is_not {
                    "NOT"
                } else {
                    ""
                };

                let relation = get_gpad_relation_of(term, annotation_detail);

                let ontology_class_id = &term_annotation.term;
                let reference_uniquename =
                    annotation_detail.reference.clone()
                    .unwrap_or_else(|| RcString::from(""));
                let evidence_type =
                    eco_evidence_from_annotation(go_eco_mappping, annotation_detail);

                let with_iter = annotation_detail.withs.iter();
                let from_iter = annotation_detail.froms.iter();
                let mut with_or_from_parts =
                    with_iter.chain(from_iter).map(|s| {
                        let s: RcString = s.clone().into();
                        if s.contains(":") {
                            s
                        } else {
                            RcString::from(&format!("{}:{}", assigned_by, s))
                        }
                    })
                    .collect::<Vec<RcString>>();
                with_or_from_parts.sort_unstable();
                let with_or_from = with_or_from_parts.join(",");

                if let Some(ref date) = annotation_detail.date {
                    let annotation_extensions =
                        make_gpad_extension_string(config, &annotation_detail.extension);
                    let line = format!("{}\t{}\t{}\t{}\t{}\t{}\t{}\t\t{}\t{}\t{}\t\n",
                                       db_object_id,
                                       not, relation, ontology_class_id,
                                       reference_uniquename, evidence_type,
                                       with_or_from, date, assigned_by,
                                       annotation_extensions);
                    gpad_writer.write_all(line.as_bytes())?;
                } else {
                    println!("WARNING while writing GAF file: \
                              date missing from annotation {}", annotation_id);
                }
            }
        }
    }

    Ok(())
}

fn make_gaf_extension_string(config: &Config, extension: &Vec<ExtPart>) -> String {
    let mut ret_string = String::new();

    for (idx, ext_part) in extension.iter().enumerate() {
        ret_string += &ext_part.rel_type_name;
        ret_string.push('(');

        if let ExtRange::Gene(ref gene_uniquename) = ext_part.ext_range {
            if gene_uniquename.contains(":") {
                ret_string += gene_uniquename;
            } else {
                ret_string += &config.database_name;
                ret_string.push(':');
                ret_string += &gene_uniquename;
            }
        } else {
            write!(ret_string, "{}", &ext_part.ext_range).unwrap();
        }
        ret_string.push(')');

        if idx < extension.len() - 1 {
            ret_string.push(',');
        }
    }

    ret_string
}

pub fn write_go_annotation_format(writer: &mut dyn io::Write, config: &Config,
                                  api_maps: &APIMaps, gene_details: &GeneDetails,
                                  cv_name: &str)
                                  -> Result<(), io::Error>
{
    if !GO_ASPECT_NAMES.contains(&cv_name) {
        return Err(io::Error::new(io::ErrorKind::InvalidInput,
                                  format!("unknown CV: {}", cv_name)));
    }

    let database_name = &config.database_name;
    let db_object_id = &gene_details.uniquename;
    let db_object_symbol =
        if let Some(ref gene_name) = gene_details.name {
            gene_name.as_str()
        } else {
            &db_object_id
        };
    let db_object_synonyms =
        gene_details.synonyms.iter().filter(|synonym| {
            synonym.synonym_type == "exact"
        })
        .map(|synonym| synonym.name.to_string())
        .collect::<Vec<String>>()
        .join("|");

    let single_letter_aspect =
        if cv_name.starts_with("b") {
            "P"
        } else {
            if cv_name.starts_with("c") {
                "C"
            } else {
                "F"
            }
        };

    if let Some(term_annotations) = gene_details.cv_annotations.get(cv_name) {
        for term_annotation in term_annotations {
            for annotation_id in &term_annotation.annotations {
                let annotation_detail = api_maps.annotation_details
                    .get(&annotation_id)
                    .expect(&format!("can't find annotation {}", annotation_id));

                let mut qualifier_parts = vec![];
                if term_annotation.is_not {
                    qualifier_parts.push("NOT");
                };

                qualifier_parts.extend(annotation_detail.qualifiers.iter()
                                       .map(|s| s.as_str()));

                let go_id = &term_annotation.term;
                let reference_uniquename =
                    annotation_detail.reference.clone()
                    .unwrap_or_else(|| RcString::from(""));
                let assigned_by =
                    if let Some(ref assigned_by) = annotation_detail.assigned_by {
                        assigned_by.as_str()
                    } else {
                        database_name.as_str()
                    };

                let with_iter = annotation_detail.withs.iter();
                let from_iter = annotation_detail.froms.iter();
                let mut with_or_from_parts =
                    with_iter.chain(from_iter).map(|s| {
                        let s: RcString = s.clone().into();
                        if s.contains(":") {
                            s
                        } else {
                            RcString::from(&format!("{}:{}", assigned_by, s))
                        }
                    })
                    .collect::<Vec<RcString>>();
                with_or_from_parts.sort_unstable();
                let with_or_from = with_or_from_parts.join(",");

                let evidence_code =
                    annotation_detail.evidence
                    .as_ref()
                    .map(|s| s.as_str())
                    .unwrap_or_else(|| "");

                let db_object_name =
                    if let Some(ref product) = gene_details.product {
                        product.as_str()
                    } else {
                        ""
                    };

                let transcript_type_str =
                    gene_details.transcripts[0].transcript_type.as_str();
                let db_object_type =
                    if transcript_type_str == "mRNA" {
                        "protein"
                    } else {
                        transcript_type_str
                    };
                let date =
                    if let Some(ref raw_date) = annotation_detail.date {
                        raw_date.replace("-", "")
                    } else {
                        println!("WARNING while writing GPAD file: \
                                  date missing from annotation {}", annotation_id);
                        continue;
                    };

                let annotation_extensions =
                    make_gaf_extension_string(config, &annotation_detail.extension);

                let gene_product_form_id =
                    if let Some(ref gene_product_form_id) = annotation_detail.gene_product_form_id
                {
                    gene_product_form_id.as_str()
                } else {
                    ""
                };

                let line = format!("{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\ttaxon:{}\t{}\t{}\t{}\t{}\n",
                                   database_name,
                                   db_object_id,
                                   db_object_symbol,
                                   qualifier_parts.join("|"),
                                   go_id,
                                   reference_uniquename,
                                   evidence_code,
                                   with_or_from,
                                   single_letter_aspect,
                                   db_object_name,
                                   db_object_synonyms,
                                   db_object_type,
                                   gene_details.taxonid,
                                   date,
                                   assigned_by,
                                   annotation_extensions,
                                   gene_product_form_id);
                writer.write_all(line.as_bytes())?;
            }
        }
    }

    Ok(())
}

<|endoftext|>


<commit_before>// Copyright 2017-2019 int08h LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//!
//! Implements the Roughenough server functionality.
//!

use hex;
use std::io::ErrorKind;
use std::net::{SocketAddr, IpAddr};
use std::process;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use std::io::Write;

use byteorder::{LittleEndian, WriteBytesExt};

use humansize::{FileSize, file_size_opts as fsopts};

use mio::net::{TcpListener, UdpSocket};
use mio::{Events, Poll, PollOpt, Ready, Token};
use mio::tcp::Shutdown;
use mio_extras::timer::Timer;

use crate::{Error, RtMessage, Tag, MIN_REQUEST_LENGTH};
use crate::config::ServerConfig;
use crate::grease::Grease;
use crate::key::{LongTermKey, OnlineKey};
use crate::kms;
use crate::merkle::MerkleTree;
use crate::stats::{AggregatedStats, PerClientStats, ServerStats, ClientStatEntry};

macro_rules! check_ctrlc {
    ($keep_running:expr) => {
        if !$keep_running.load(Ordering::Acquire) {
            warn!("Ctrl-C caught, exiting...");
            return true;
        }
    };
}

// mio event registrations
const MESSAGE: Token = Token(0);
const STATUS: Token = Token(1);
const HEALTH_CHECK: Token = Token(2);

// Canned response to health check request
const HTTP_RESPONSE: &str = "HTTP/1.1 200 OK\nContent-Length: 0\nConnection: close\n\n";

/// The main Roughenough server instance.
///
/// The [ServerConfig](../config/trait.ServerConfig.html) trait specifies the required and optional
/// parameters available for configuring a Roughenoguh server instance.
///
/// Implementations of `ServerConfig` obtain configurations from different back-end sources
/// such as files or environment variables.
///
/// See [the config module](../config/index.html) for more information.
///
pub struct Server {
    config: Box<dyn ServerConfig>,
    online_key: OnlineKey,
    cert_bytes: Vec<u8>,

    socket: UdpSocket,
    health_listener: Option<TcpListener>,
    keep_running: Arc<AtomicBool>,
    poll_duration: Option<Duration>,
    grease: Grease,
    timer: Timer<()>,
    poll: Poll,
    merkle: MerkleTree,
    requests: Vec<(Vec<u8>, SocketAddr)>,
    buf: [u8; 65_536],

    public_key: String,

    stats: Box<dyn ServerStats>,

    // Used to send requests to ourselves in fuzzing mode
    #[cfg(fuzzing)]
    fake_client_socket: UdpSocket,
}

impl Server {
    ///
    /// Create a new server instance from the provided
    /// [`ServerConfig`](../config/trait.ServerConfig.html) trait object instance.
    ///
    pub fn new(config: Box<dyn ServerConfig>) -> Server {
        let online_key = OnlineKey::new();
        let public_key: String;

        let cert_bytes = {
            let seed = match kms::load_seed(config.as_ref()) {
                Ok(seed) => seed,
                Err(e) => {
                    error!("Failed to load seed: {:#?}", e);
                    process::exit(1);
                }
            };
            let mut long_term_key = LongTermKey::new(&seed);
            public_key = hex::encode(long_term_key.public_key());

            long_term_key.make_cert(&online_key).encode().unwrap()
        };

        let keep_running = Arc::new(AtomicBool::new(true));

        let sock_addr = config.udp_socket_addr().expect("udp sock addr");
        let socket = UdpSocket::bind(&sock_addr).expect("failed to bind to socket");

        let poll_duration = Some(Duration::from_millis(100));

        let mut timer: Timer<()> = Timer::default();
        timer.set_timeout(config.status_interval(), ());

        let poll = Poll::new().unwrap();
        poll.register(&socket, MESSAGE, Ready::readable(), PollOpt::edge())
            .unwrap();
        poll.register(&timer, STATUS, Ready::readable(), PollOpt::edge())
            .unwrap();

        let health_listener = if let Some(hc_port) = config.health_check_port() {
            let hc_sock_addr: SocketAddr = format!("{}:{}", config.interface(), hc_port)
                .parse()
                .unwrap();

            let tcp_listener = TcpListener::bind(&hc_sock_addr)
                .expect("failed to bind TCP listener for health check");

            poll.register(&tcp_listener, HEALTH_CHECK, Ready::readable(), PollOpt::edge())
                .unwrap();

            Some(tcp_listener)
        } else {
            None
        };

        let stats: Box<dyn ServerStats> = if config.client_stats_enabled() {
            Box::new(PerClientStats::new())
        } else {
            Box::new(AggregatedStats::new())
        };

        let grease = Grease::new(config.fault_percentage());
        let merkle = MerkleTree::new();
        let requests = Vec::with_capacity(config.batch_size() as usize);

        Server {
            config,
            online_key,
            cert_bytes,

            socket,
            health_listener,

            keep_running,
            poll_duration,
            grease,
            timer,
            poll,
            merkle,
            requests,
            buf: [0u8; 65_536],

            public_key,

            stats,

            #[cfg(fuzzing)]
            fake_client_socket: UdpSocket::bind(&"127.0.0.1:0".parse().unwrap()).unwrap(),
        }
    }

    /// Returns a reference to the server's long-term public key
    pub fn get_public_key(&self) -> &str {
        &self.public_key
    }

    /// Returns a reference to the server's on-line (delegated) key
    pub fn get_online_key(&self) -> &OnlineKey {
        &self.online_key
    }

    /// Returns a reference to the `ServerConfig` this server was configured with
    pub fn get_config(&self) -> &dyn ServerConfig {
        self.config.as_ref()
    }

    /// Returns a reference counted pointer the this server's `keep_running` value.
    pub fn get_keep_running(&self) -> Arc<AtomicBool> {
        self.keep_running.clone()
    }

    #[cfg(fuzzing)]
    pub fn send_to_self(&mut self, data: &[u8]) {
        let res = self
            .fake_client_socket
            .send_to(data, &self.socket.local_addr().unwrap());
        info!("Sent to self: {:?}", res);
    }

    /// The main processing function for incoming connections. This method should be
    /// called repeatedly in a loop to process requests. It returns 'true' when the
    /// server has shutdown (due to keep_running being set to 'false').
    ///
    pub fn process_events(&mut self, events: &mut Events) -> bool {
        self.poll
            .poll(events, self.poll_duration)
            .expect("server event poll failed; cannot recover");

        for msg in events.iter() {
            match msg.token() {
                MESSAGE => {
                    loop {
                        check_ctrlc!(self.keep_running);

                        self.merkle.reset();
                        self.requests.clear();

                        let socket_now_empty = self.collect_requests();

                        if self.requests.is_empty() {
                            break;
                        }

                        self.send_responses();

                        if socket_now_empty {
                            break;
                        }
                    }
                }

                HEALTH_CHECK => {
                    let listener = self.health_listener.as_ref().unwrap();

                    match listener.accept() {
                        Ok((ref mut stream, src_addr)) => {
                            info!("health check from {}", src_addr);
                            self.stats.add_health_check(&src_addr.ip());

                            match stream.write(HTTP_RESPONSE.as_bytes()) {
                                Ok(_) => (),
                                Err(e) => warn!("error writing health check {}", e),
                            };

                            match stream.shutdown(Shutdown::Both) {
                                Ok(_) => (),
                                Err(e) => warn!("error in health check socket shutdown {}", e),
                            }
                        }
                        Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                            debug!("blocking in TCP health check");
                        }
                        Err(e) => {
                            warn!("unexpected health check error {}", e);
                        }
                    }
                }

                STATUS => {
                    let mut vec: Vec<(&IpAddr, &ClientStatEntry)> = self.stats.iter().collect();
                    vec.sort_by(|lhs, rhs| lhs.1.valid_requests.cmp(&rhs.1.valid_requests));

                    for (addr, counts) in vec {
                        info!(
                            "{:16}: {} valid, {} invalid requests; {} responses ({} sent)",
                            format!("{}", addr), counts.valid_requests, counts.invalid_requests,
                            counts.responses_sent,
                            counts.bytes_sent.file_size(fsopts::BINARY).unwrap()
                        );
                    }

                    info!(
                        "Totals: {} unique clients; {} valid, {} invalid requests; {} responses ({} sent)",
                        self.stats.total_unique_clients(),
                        self.stats.total_valid_requests(), self.stats.total_invalid_requests(),
                        self.stats.total_responses_sent(),
                        self.stats.total_bytes_sent().file_size(fsopts::BINARY).unwrap()
                    );

                    self.timer.set_timeout(self.config.status_interval(), ());
                }

                _ => unreachable!(),
            }
        }
        false
    }

    // Read and process client requests from socket until empty or 'batch_size' number of
    // requests have been read.
    fn collect_requests(&mut self) -> bool {
        for i in 0..self.config.batch_size() {
            match self.socket.recv_from(&mut self.buf) {
                Ok((num_bytes, src_addr)) => {
                    match self.nonce_from_request(&self.buf, num_bytes) {
                        Ok(nonce) => {
                            self.stats.add_valid_request(&src_addr.ip());
                            self.requests.push((Vec::from(nonce), src_addr));
                            self.merkle.push_leaf(nonce);
                        }
                        Err(e) => {
                            self.stats.add_invalid_request(&src_addr.ip());

                            info!(
                                "Invalid request: '{:?}' ({} bytes) from {} (#{} in batch, resp #{})",
                                e, num_bytes, src_addr, i,
                                self.stats.total_responses_sent() + u64::from(i)
                            );
                        }
                    }
                }
                Err(e) => match e.kind() {
                    ErrorKind::WouldBlock => {
                        return true;
                    }
                    _ => {
                        error!("Error receiving from socket: {:?}: {:?}", e.kind(), e);
                        return false;
                    }
                },
            };
        }

        false
    }

    // extract the client's nonce from its request
    fn nonce_from_request<'a>(&self, buf: &'a [u8], num_bytes: usize) -> Result<&'a [u8], Error> {
        if num_bytes < MIN_REQUEST_LENGTH as usize {
            return Err(Error::RequestTooShort);
        }

        let tag_count = &buf[..4];
        let expected_nonc = &buf[8..12];
        let expected_pad = &buf[12..16];

        let tag_count_is_2 = tag_count == [0x02, 0x00, 0x00, 0x00];
        let tag1_is_nonc = expected_nonc == Tag::NONC.wire_value();
        let tag2_is_pad = expected_pad == Tag::PAD.wire_value();

        if tag_count_is_2 && tag1_is_nonc && tag2_is_pad {
            Ok(&buf[0x10..0x50])
        } else {
            Err(Error::InvalidRequest)
        }
    }

    fn send_responses(&mut self) {
        let merkle_root = self.merkle.compute_root();

        // The SREP tag is identical for each response
        let srep = self.online_key.make_srep(SystemTime::now(), &merkle_root);

        for (i, &(ref nonce, ref src_addr)) in self.requests.iter().enumerate() {
            let paths = self.merkle.get_paths(i);
            let resp_msg = {
                let r = self.make_response(&srep, &self.cert_bytes, &paths, i as u32);
                if self.grease.should_add_error() { self.grease.add_errors(&r) } else { r }
            };
            let resp_bytes = resp_msg.encode().unwrap();

            let bytes_sent = self
                .socket
                .send_to(&resp_bytes, &src_addr)
                .expect("send_to failed");

            self.stats.add_response(&src_addr.ip(), bytes_sent);

            info!(
                "Responded {} bytes to {} for '{}..' (#{} in batch, resp #{})",
                bytes_sent,
                src_addr,
                hex::encode(&nonce[0..4]),
                i + 1,
                self.stats.total_responses_sent()
            );
        }
    }

    fn make_response(&self, srep: &RtMessage, cert_bytes: &[u8], path: &[u8], idx: u32) -> RtMessage {
        let mut index = [0; 4];
        (&mut index as &mut [u8])
            .write_u32::<LittleEndian>(idx)
            .unwrap();

        let sig_bytes = srep.get_field(Tag::SIG).unwrap();
        let srep_bytes = srep.get_field(Tag::SREP).unwrap();

        let mut response = RtMessage::new(5);
        response.add_field(Tag::SIG, sig_bytes).unwrap();
        response.add_field(Tag::PATH, path).unwrap();
        response.add_field(Tag::SREP, srep_bytes).unwrap();
        response.add_field(Tag::CERT, cert_bytes).unwrap();
        response.add_field(Tag::INDX, &index).unwrap();

        response
    }
}
<commit_msg>refactor: extract health check and status update methods out of event loop<commit_after>// Copyright 2017-2019 int08h LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//!
//! Implements the Roughenough server functionality.
//!

use hex;
use std::io::ErrorKind;
use std::net::{SocketAddr, IpAddr};
use std::process;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use std::io::Write;

use byteorder::{LittleEndian, WriteBytesExt};

use humansize::{FileSize, file_size_opts as fsopts};

use mio::net::{TcpListener, UdpSocket};
use mio::{Events, Poll, PollOpt, Ready, Token};
use mio::tcp::Shutdown;
use mio_extras::timer::Timer;

use crate::{Error, RtMessage, Tag, MIN_REQUEST_LENGTH};
use crate::config::ServerConfig;
use crate::grease::Grease;
use crate::key::{LongTermKey, OnlineKey};
use crate::kms;
use crate::merkle::MerkleTree;
use crate::stats::{AggregatedStats, PerClientStats, ServerStats, ClientStatEntry};

macro_rules! check_ctrlc {
    ($keep_running:expr) => {
        if !$keep_running.load(Ordering::Acquire) {
            warn!("Ctrl-C caught, exiting...");
            return true;
        }
    };
}

// mio event registrations
const EVT_MESSAGE: Token = Token(0);
const EVT_STATUS_UPDATE: Token = Token(1);
const EVT_HEALTH_CHECK: Token = Token(2);

// Canned response to health check request
const HTTP_RESPONSE: &str = "HTTP/1.1 200 OK\nContent-Length: 0\nConnection: close\n\n";

/// The main Roughenough server instance.
///
/// The [ServerConfig](../config/trait.ServerConfig.html) trait specifies the required and optional
/// parameters available for configuring a Roughenoguh server instance.
///
/// Implementations of `ServerConfig` obtain configurations from different back-end sources
/// such as files or environment variables.
///
/// See [the config module](../config/index.html) for more information.
///
pub struct Server {
    config: Box<dyn ServerConfig>,
    online_key: OnlineKey,
    cert_bytes: Vec<u8>,

    socket: UdpSocket,
    health_listener: Option<TcpListener>,
    keep_running: Arc<AtomicBool>,
    poll_duration: Option<Duration>,
    grease: Grease,
    timer: Timer<()>,
    poll: Poll,
    merkle: MerkleTree,
    requests: Vec<(Vec<u8>, SocketAddr)>,
    buf: [u8; 65_536],

    public_key: String,

    stats: Box<dyn ServerStats>,

    // Used to send requests to ourselves in fuzzing mode
    #[cfg(fuzzing)]
    fake_client_socket: UdpSocket,
}

impl Server {
    ///
    /// Create a new server instance from the provided
    /// [`ServerConfig`](../config/trait.ServerConfig.html) trait object instance.
    ///
    pub fn new(config: Box<dyn ServerConfig>) -> Server {
        let online_key = OnlineKey::new();
        let public_key: String;

        let cert_bytes = {
            let seed = match kms::load_seed(config.as_ref()) {
                Ok(seed) => seed,
                Err(e) => {
                    error!("Failed to load seed: {:#?}", e);
                    process::exit(1);
                }
            };
            let mut long_term_key = LongTermKey::new(&seed);
            public_key = hex::encode(long_term_key.public_key());

            long_term_key.make_cert(&online_key).encode().unwrap()
        };

        let keep_running = Arc::new(AtomicBool::new(true));

        let sock_addr = config.udp_socket_addr().expect("udp sock addr");
        let socket = UdpSocket::bind(&sock_addr).expect("failed to bind to socket");

        let poll_duration = Some(Duration::from_millis(100));

        let mut timer: Timer<()> = Timer::default();
        timer.set_timeout(config.status_interval(), ());

        let poll = Poll::new().unwrap();
        poll.register(&socket, EVT_MESSAGE, Ready::readable(), PollOpt::edge())
            .unwrap();
        poll.register(&timer, EVT_STATUS_UPDATE, Ready::readable(), PollOpt::edge())
            .unwrap();

        let health_listener = if let Some(hc_port) = config.health_check_port() {
            let hc_sock_addr: SocketAddr = format!("{}:{}", config.interface(), hc_port)
                .parse()
                .unwrap();

            let tcp_listener = TcpListener::bind(&hc_sock_addr)
                .expect("failed to bind TCP listener for health check");

            poll.register(&tcp_listener, EVT_HEALTH_CHECK, Ready::readable(), PollOpt::edge())
                .unwrap();

            Some(tcp_listener)
        } else {
            None
        };

        let stats: Box<dyn ServerStats> = if config.client_stats_enabled() {
            Box::new(PerClientStats::new())
        } else {
            Box::new(AggregatedStats::new())
        };

        let grease = Grease::new(config.fault_percentage());
        let merkle = MerkleTree::new();
        let requests = Vec::with_capacity(config.batch_size() as usize);

        Server {
            config,
            online_key,
            cert_bytes,

            socket,
            health_listener,

            keep_running,
            poll_duration,
            grease,
            timer,
            poll,
            merkle,
            requests,
            buf: [0u8; 65_536],

            public_key,

            stats,

            #[cfg(fuzzing)]
            fake_client_socket: UdpSocket::bind(&"127.0.0.1:0".parse().unwrap()).unwrap(),
        }
    }

    /// Returns a reference to the server's long-term public key
    pub fn get_public_key(&self) -> &str {
        &self.public_key
    }

    /// Returns a reference to the server's on-line (delegated) key
    pub fn get_online_key(&self) -> &OnlineKey {
        &self.online_key
    }

    /// Returns a reference to the `ServerConfig` this server was configured with
    pub fn get_config(&self) -> &dyn ServerConfig {
        self.config.as_ref()
    }

    /// Returns a reference counted pointer the this server's `keep_running` value.
    pub fn get_keep_running(&self) -> Arc<AtomicBool> {
        self.keep_running.clone()
    }

    #[cfg(fuzzing)]
    pub fn send_to_self(&mut self, data: &[u8]) {
        let res = self
            .fake_client_socket
            .send_to(data, &self.socket.local_addr().unwrap());
        info!("Sent to self: {:?}", res);
    }

    /// The main processing function for incoming connections. This method should be
    /// called repeatedly in a loop to process requests. It returns 'true' when the
    /// server has shutdown (due to keep_running being set to 'false').
    ///
    pub fn process_events(&mut self, events: &mut Events) -> bool {
        self.poll
            .poll(events, self.poll_duration)
            .expect("server event poll failed; cannot recover");

        for msg in events.iter() {
            match msg.token() {
                EVT_MESSAGE => {
                    loop {
                        check_ctrlc!(self.keep_running);

                        self.merkle.reset();
                        self.requests.clear();

                        let socket_now_empty = self.collect_requests();

                        if self.requests.is_empty() {
                            break;
                        }

                        self.send_responses();

                        if socket_now_empty {
                            break;
                        }
                    }
                },
                EVT_HEALTH_CHECK => self.handle_health_check(),
                EVT_STATUS_UPDATE => self.handle_status_update(),
                _ => unreachable!(),
            }
        }
        false
    }

    // Read and process client requests from socket until empty or 'batch_size' number of
    // requests have been read.
    fn collect_requests(&mut self) -> bool {
        for i in 0..self.config.batch_size() {
            match self.socket.recv_from(&mut self.buf) {
                Ok((num_bytes, src_addr)) => {
                    match self.nonce_from_request(&self.buf, num_bytes) {
                        Ok(nonce) => {
                            self.stats.add_valid_request(&src_addr.ip());
                            self.requests.push((Vec::from(nonce), src_addr));
                            self.merkle.push_leaf(nonce);
                        }
                        Err(e) => {
                            self.stats.add_invalid_request(&src_addr.ip());

                            info!(
                                "Invalid request: '{:?}' ({} bytes) from {} (#{} in batch, resp #{})",
                                e, num_bytes, src_addr, i,
                                self.stats.total_responses_sent() + u64::from(i)
                            );
                        }
                    }
                }
                Err(e) => match e.kind() {
                    ErrorKind::WouldBlock => {
                        return true;
                    }
                    _ => {
                        error!("Error receiving from socket: {:?}: {:?}", e.kind(), e);
                        return false;
                    }
                },
            };
        }

        false
    }

    // extract the client's nonce from its request
    fn nonce_from_request<'a>(&self, buf: &'a [u8], num_bytes: usize) -> Result<&'a [u8], Error> {
        if num_bytes < MIN_REQUEST_LENGTH as usize {
            return Err(Error::RequestTooShort);
        }

        let tag_count = &buf[..4];
        let expected_nonc = &buf[8..12];
        let expected_pad = &buf[12..16];

        let tag_count_is_2 = tag_count == [0x02, 0x00, 0x00, 0x00];
        let tag1_is_nonc = expected_nonc == Tag::NONC.wire_value();
        let tag2_is_pad = expected_pad == Tag::PAD.wire_value();

        if tag_count_is_2 && tag1_is_nonc && tag2_is_pad {
            Ok(&buf[0x10..0x50])
        } else {
            Err(Error::InvalidRequest)
        }
    }

    fn send_responses(&mut self) {
        let merkle_root = self.merkle.compute_root();

        // The SREP tag is identical for each response
        let srep = self.online_key.make_srep(SystemTime::now(), &merkle_root);

        for (i, &(ref nonce, ref src_addr)) in self.requests.iter().enumerate() {
            let paths = self.merkle.get_paths(i);
            let resp_msg = {
                let r = self.make_response(&srep, &self.cert_bytes, &paths, i as u32);
                if self.grease.should_add_error() { self.grease.add_errors(&r) } else { r }
            };
            let resp_bytes = resp_msg.encode().unwrap();

            let bytes_sent = self
                .socket
                .send_to(&resp_bytes, &src_addr)
                .expect("send_to failed");

            self.stats.add_response(&src_addr.ip(), bytes_sent);

            info!(
                "Responded {} bytes to {} for '{}..' (#{} in batch, resp #{})",
                bytes_sent,
                src_addr,
                hex::encode(&nonce[0..4]),
                i + 1,
                self.stats.total_responses_sent()
            );
        }
    }

    fn make_response(&self, srep: &RtMessage, cert_bytes: &[u8], path: &[u8], idx: u32) -> RtMessage {
        let mut index = [0; 4];
        (&mut index as &mut [u8])
            .write_u32::<LittleEndian>(idx)
            .unwrap();

        let sig_bytes = srep.get_field(Tag::SIG).unwrap();
        let srep_bytes = srep.get_field(Tag::SREP).unwrap();

        let mut response = RtMessage::new(5);
        response.add_field(Tag::SIG, sig_bytes).unwrap();
        response.add_field(Tag::PATH, path).unwrap();
        response.add_field(Tag::SREP, srep_bytes).unwrap();
        response.add_field(Tag::CERT, cert_bytes).unwrap();
        response.add_field(Tag::INDX, &index).unwrap();

        response
    }

    fn handle_health_check(&mut self) {
        let listener = self.health_listener.as_ref().unwrap();
        match listener.accept() {
            Ok((ref mut stream, src_addr)) => {
                info!("health check from {}", src_addr);
                self.stats.add_health_check(&src_addr.ip());

                match stream.write(HTTP_RESPONSE.as_bytes()) {
                    Ok(_) => (),
                    Err(e) => warn!("error writing health check {}", e),
                };

                match stream.shutdown(Shutdown::Both) {
                    Ok(_) => (),
                    Err(e) => warn!("error in health check socket shutdown {}", e),
                }
            }
            Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                debug!("blocking in TCP health check");
            }
            Err(e) => {
                warn!("unexpected health check error {}", e);
            }
        }
    }

    fn handle_status_update(&mut self) {
        let mut vec: Vec<(&IpAddr, &ClientStatEntry)> = self.stats.iter().collect();
        vec.sort_by(|lhs, rhs| lhs.1.valid_requests.cmp(&rhs.1.valid_requests));

        for (addr, counts) in vec {
            info!(
                "{:16}: {} valid, {} invalid requests; {} responses ({} sent)",
                format!("{}", addr), counts.valid_requests, counts.invalid_requests,
                counts.responses_sent,
                counts.bytes_sent.file_size(fsopts::BINARY).unwrap()
            );
        }

        info!(
            "Totals: {} unique clients; {} valid, {} invalid requests; {} responses ({} sent)",
            self.stats.total_unique_clients(),
            self.stats.total_valid_requests(), self.stats.total_invalid_requests(),
            self.stats.total_responses_sent(),
            self.stats.total_bytes_sent().file_size(fsopts::BINARY).unwrap()
        );

        self.timer.set_timeout(self.config.status_interval(), ());
    }
}
<|endoftext|>


<commit_before>//! Models the static structure of the database which Ardite is serving. All
//! Ardite services need some static representation of the database which
//! contains just enough information to get started. The types in this module
//! are that bootstrap static structure.

/// Represents a collection of documents. May contain whatever necessary data
/// to *locate* the collection. These collections must be simple enough to do
/// easy equality checks.
pub trait Collection {
  /// Get the public name of the collection. In a PostgreSQL database, this
  /// would be the table name without the schema name.
  fn get_name() -> String;
}
<commit_msg>Add structure struct back in<commit_after>//! Models the static structure of the database which Ardite is serving. All
//! Ardite services need some static representation of the database which
//! contains just enough information to get started. The types in this module
//! are that bootstrap static structure.

/// Represents the entire database structure.
pub struct Structure<C: Collection> {
  /// All of the collections in the database. May or may not be accessible by
  /// Ardite services.
  collections: Vec<C>
}

/// Represents a collection of documents. May contain whatever necessary data
/// to *locate* the collection. These collections must be simple enough to do
/// easy equality checks.
pub trait Collection {
  /// Get the public name of the collection. In a PostgreSQL database, this
  /// would be the table name without the schema name.
  fn get_name() -> String;
}
<|endoftext|>


<commit_before>use {Errno, Result};
use libc::c_int;
use std::os::unix::io::RawFd;

mod ffi {
    use libc::{c_int};
    use super::EpollEvent;

    extern {
        pub fn epoll_create(size: c_int) -> c_int;
        pub fn epoll_create1(flags: c_int) -> c_int;
        pub fn epoll_ctl(epfd: c_int, op: c_int, fd: c_int, event: *const EpollEvent) -> c_int;
        pub fn epoll_wait(epfd: c_int, events: *mut EpollEvent, max_events: c_int, timeout: c_int) -> c_int;
    }
}

bitflags!(
    flags EpollFdFlag: c_int {
        const EPOLL_NONBLOCK = 0x800,
        const EPOLL_CLOEXEC = 0x80000
    }
);

bitflags!(
    #[repr(C)]
    flags EpollEventKind: u32 {
        const EPOLLIN = 0x001,
        const EPOLLPRI = 0x002,
        const EPOLLOUT = 0x004,
        const EPOLLRDNORM = 0x040,
        const EPOLLRDBAND = 0x080,
        const EPOLLWRNORM = 0x100,
        const EPOLLWRBAND = 0x200,
        const EPOLLMSG = 0x400,
        const EPOLLERR = 0x008,
        const EPOLLHUP = 0x010,
        const EPOLLRDHUP = 0x2000,
        const EPOLLEXCLUSIVE = 1 << 28,
        const EPOLLWAKEUP = 1 << 29,
        const EPOLLONESHOT = 1 << 30,
        const EPOLLET = 1 << 31
    }
);

#[derive(Clone, Copy)]
#[repr(C)]
pub enum EpollOp {
    EpollCtlAdd = 1,
    EpollCtlDel = 2,
    EpollCtlMod = 3
}

#[cfg(not(target_arch = "x86_64"))]
#[derive(Clone, Copy)]
#[repr(C)]
pub struct EpollEvent {
    pub events: EpollEventKind,
    pub data: u64
}

#[cfg(target_arch = "x86_64")]
#[derive(Clone, Copy)]
#[repr(C, packed)]
pub struct EpollEvent {
    pub events: EpollEventKind,
    pub data: u64
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[test]
fn test_epoll_event_size() {
    use std::mem::size_of;
    assert_eq!(size_of::<EpollEvent>(), 12);
}

#[cfg(target_arch = "arm")]
#[test]
fn test_epoll_event_size() {
    use std::mem::size_of;
    assert_eq!(size_of::<EpollEvent>(), 16);
}

#[inline]
pub fn epoll_create() -> Result<RawFd> {
    let res = unsafe { ffi::epoll_create(1024) };

    Errno::result(res)
}

#[inline]
pub fn epoll_create1(flags: EpollFdFlag) -> Result<RawFd> {
    let res = unsafe { ffi::epoll_create1(flags.bits()) };

    Errno::result(res)
}

#[inline]
pub fn epoll_ctl(epfd: RawFd, op: EpollOp, fd: RawFd, event: &EpollEvent) -> Result<()> {
    let res = unsafe { ffi::epoll_ctl(epfd, op as c_int, fd, event as *const EpollEvent) };

    Errno::result(res).map(drop)
}

#[inline]
pub fn epoll_wait(epfd: RawFd, events: &mut [EpollEvent], timeout_ms: isize) -> Result<usize> {
    let res = unsafe {
        ffi::epoll_wait(epfd, events.as_mut_ptr(), events.len() as c_int, timeout_ms as c_int)
    };

    Errno::result(res).map(|r| r as usize)
}
<commit_msg>Use libc's declarations<commit_after>use {Errno, Result};
use libc::{self, c_int};
use std::os::unix::io::RawFd;

bitflags!(
    #[repr(C)]
    flags EpollEventKind: u32 {
        const EPOLLIN = 0x001,
        const EPOLLPRI = 0x002,
        const EPOLLOUT = 0x004,
        const EPOLLRDNORM = 0x040,
        const EPOLLRDBAND = 0x080,
        const EPOLLWRNORM = 0x100,
        const EPOLLWRBAND = 0x200,
        const EPOLLMSG = 0x400,
        const EPOLLERR = 0x008,
        const EPOLLHUP = 0x010,
        const EPOLLRDHUP = 0x2000,
        const EPOLLEXCLUSIVE = 1 << 28,
        const EPOLLWAKEUP = 1 << 29,
        const EPOLLONESHOT = 1 << 30,
        const EPOLLET = 1 << 31
    }
);

#[derive(Clone, Copy)]
#[repr(C)]
pub enum EpollOp {
    EpollCtlAdd = 1,
    EpollCtlDel = 2,
    EpollCtlMod = 3
}

#[inline]
pub fn epoll_create() -> Result<RawFd> {
    let res = unsafe { libc::epoll_create(1024) };

    Errno::result(res)
}

#[inline]
pub fn epoll_create1(flags: c_int) -> Result<RawFd> {
    let res = unsafe { libc::epoll_create1(flags) };

    Errno::result(res)
}

#[inline]
pub fn epoll_ctl(epfd: RawFd, op: EpollOp, fd: RawFd, event: &mut libc::epoll_event) -> Result<()> {
    let res = unsafe { libc::epoll_ctl(epfd, op as c_int, fd, event as *mut libc::epoll_event) };

    Errno::result(res).map(drop)
}

#[inline]
pub fn epoll_wait(epfd: RawFd, events: &mut [libc::epoll_event], timeout_ms: isize) -> Result<usize> {
    let res = unsafe {
        libc::epoll_wait(epfd, events.as_mut_ptr(), events.len() as c_int, timeout_ms as c_int)
    };

    Errno::result(res).map(|r| r as usize)
}
<|endoftext|>