<commit_before>use byteorder::{LittleEndian, ReadBytesExt};
use crate::internal::prelude::*;
use opus::{
    Channels,
    Decoder as OpusDecoder,
    Result as OpusResult,
};
use parking_lot::Mutex;
use serde_json;
use std::{
    ffi::OsStr,
    fs::File,
    io::{BufReader, ErrorKind as IoErrorKind, Read, Result as IoResult},
    process::{Child, Command, Stdio},
    result::Result as StdResult,
    sync::Arc,
};
use super::{AudioSource, AudioType, DcaError, DcaMetadata, VoiceError, audio};
use log::{debug, warn};

struct ChildContainer(Child);

impl Read for ChildContainer {
    fn read(&mut self, buffer: &mut [u8]) -> IoResult<usize> {
        self.0.stdout.as_mut().unwrap().read(buffer)
    }
}

impl Drop for ChildContainer {
    fn drop (&mut self) {
        if let Err(e) = self.0.kill() {
            debug!("[Voice] Error awaiting child process: {:?}", e);
        }
    }
}

// Since each audio item needs its own decoder, we need to
// work around the fact that OpusDecoders aint sendable.
struct SendDecoder(OpusDecoder);

impl SendDecoder {
    fn decode_float(&mut self, input: &[u8], output: &mut [f32], fec: bool) -> OpusResult<usize> {
        let &mut SendDecoder(ref mut sd) = self;
        sd.decode_float(input, output, fec)
    }
}

unsafe impl Send for SendDecoder {}

struct InputSource<R: Read + Send + 'static> {
    stereo: bool,
    reader: R,
    kind: AudioType,
    decoder: Option<Arc<Mutex<SendDecoder>>>,
}

impl<R: Read + Send> AudioSource for InputSource<R> {
    fn is_stereo(&mut self) -> bool { self.stereo }

    fn get_type(&self) -> AudioType { self.kind }

    fn read_pcm_frame(&mut self, buffer: &mut [i16]) -> Option<usize> {
        for (i, v) in buffer.iter_mut().enumerate() {
            *v = match self.reader.read_i16::<LittleEndian>() {
                Ok(v) => v,
                Err(ref e) => {
                    return if e.kind() == IoErrorKind::UnexpectedEof {
                        Some(i)
                    } else {
                        None
                    }
                },
            }
        }

        Some(buffer.len())
    }

    fn read_opus_frame(&mut self) -> Option<Vec<u8>> {
        match self.reader.read_i16::<LittleEndian>() {
            Ok(size) => {
                if size <= 0 {
                    warn!("Invalid opus frame size: {}", size);
                    return None;
                }

                let mut frame = Vec::with_capacity(size as usize);

                {
                    let reader = self.reader.by_ref();

                    if reader.take(size as u64).read_to_end(&mut frame).is_err() {
                        return None;
                    }
                }

                Some(frame)
            },
            Err(ref e) => if e.kind() == IoErrorKind::UnexpectedEof {
                Some(Vec::new())
            } else {
                None
            },
        }
    }

    fn decode_and_add_opus_frame(&mut self, float_buffer: &mut [f32; 1920], volume: f32) -> Option<usize> {
        let decoder_lock = self.decoder.as_mut()?.clone();
        let frame = self.read_opus_frame()?;
        let mut local_buf = [0f32; 960 * 2];

        let count = {
            let mut decoder = decoder_lock.lock();

            decoder.decode_float(frame.as_slice(), &mut local_buf, false).ok()?
        };

        for i in 0..1920 {
            float_buffer[i] += local_buf[i] * volume;
        }

        Some(count)
    }
}

/// Opens an audio file through `ffmpeg` and creates an audio source.
pub fn ffmpeg<P: AsRef<OsStr>>(path: P) -> Result<Box<dyn AudioSource>> {
    _ffmpeg(path.as_ref())
}

fn _ffmpeg(path: &OsStr) -> Result<Box<dyn AudioSource>> {
    // Will fail if the path is not to a file on the fs. Likely a YouTube URI.
    let is_stereo = is_stereo(path).unwrap_or(false);
    let stereo_val = if is_stereo { "2" } else { "1" };

    _ffmpeg_optioned(path, &[
        "-f",
        "s16le",
        "-ac",
        stereo_val,
        "-ar",
        "48000",
        "-acodec",
        "pcm_s16le",
        "-",
    ], Some(is_stereo))
}

/// Opens an audio file through `ffmpeg` and creates an audio source, with
/// user-specified arguments to pass to ffmpeg.
///
/// Note that this does _not_ build on the arguments passed by the [`ffmpeg`]
/// function.
///
/// # Examples
///
/// Pass options to create a custom ffmpeg streamer:
///
/// ```rust,no_run
/// use serenity::voice;
///
/// let stereo_val = "2";
///
/// let streamer = voice::ffmpeg_optioned("./some_file.mp3", &[
///     "-f",
///     "s16le",
///     "-ac",
///     stereo_val,
///     "-ar",
///     "48000",
///     "-acodec",
///     "pcm_s16le",
///     "-",
/// ]);
pub fn ffmpeg_optioned<P: AsRef<OsStr>>(
    path: P,
    args: &[&str],
) -> Result<Box<dyn AudioSource>> {
    _ffmpeg_optioned(path.as_ref(), args)
}

fn _ffmpeg_optioned(path: &OsStr, args: &[&str]) -> Result<Box<dyn AudioSource>> {
    let is_stereo = is_stereo(path).unwrap_or(false);

    let command = Command::new("ffmpeg")
        .arg("-i")
        .arg(path)
        .args(args)
        .stderr(Stdio::null())
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .spawn()?;

    Ok(pcm(is_stereo, ChildContainer(command)))
}

/// Creates a streamed audio source from a DCA file.
/// Currently only accepts the DCA1 format.
pub fn dca<P: AsRef<OsStr>>(path: P) -> StdResult<Box<dyn AudioSource>, DcaError> {
    _dca(path.as_ref())
}

fn _dca(path: &OsStr) -> StdResult<Box<dyn AudioSource>, DcaError> {
    let file = File::open(path).map_err(DcaError::IoError)?;

    let mut reader = BufReader::new(file);

    let mut header = [0u8; 4];

    // Read in the magic number to verify it's a DCA file.
    reader.read_exact(&mut header).map_err(DcaError::IoError)?;

    if header != b"DCA1"[..] {
        return Err(DcaError::InvalidHeader);
    }

    reader.read_exact(&mut header).map_err(DcaError::IoError)?;

    let size = (&header[..]).read_i32::<LittleEndian>().unwrap();

    // Sanity check
    if size < 2 {
        return Err(DcaError::InvalidSize(size));
    }

    let mut raw_json = Vec::with_capacity(size as usize);

    {
        let json_reader = reader.by_ref();
        json_reader
            .take(size as u64)
            .read_to_end(&mut raw_json)
            .map_err(DcaError::IoError)?;
    }

    let metadata = serde_json::from_slice::<DcaMetadata>(raw_json.as_slice())
        .map_err(DcaError::InvalidMetadata)?;

    Ok(opus(metadata.is_stereo(), reader))
}

/// Creates an Opus audio source. This makes certain assumptions: namely, that the input stream
/// is composed ONLY of opus frames of the variety that Discord expects.
///
/// If you want to decode a `.opus` file, use [`ffmpeg`]
///
/// [`ffmpeg`]: fn.ffmpeg.html
pub fn opus<R: Read + Send + 'static>(is_stereo: bool, reader: R) -> Box<dyn AudioSource> {
    Box::new(InputSource {
        stereo: is_stereo,
        reader,
        kind: AudioType::Opus,
        decoder: Some(
            Arc::new(Mutex::new(
                // We always want to decode *to* stereo, for mixing reasons.
                SendDecoder(OpusDecoder::new(audio::SAMPLE_RATE, Channels::Stereo).unwrap())
            ))
        ),
    })
}

/// Creates a PCM audio source.
pub fn pcm<R: Read + Send + 'static>(is_stereo: bool, reader: R) -> Box<dyn AudioSource> {
    Box::new(InputSource {
        stereo: is_stereo,
        reader,
        kind: AudioType::Pcm,
        decoder: None,
    })
}

/// Creates a streamed audio source with `youtube-dl` and `ffmpeg`.
pub fn ytdl(uri: &str) -> Result<Box<dyn AudioSource>> {
    let args = [
        "-f",
        "webm[abr>0]/bestaudio/best",
        "--no-playlist",
        "--print-json",
        "--skip-download",
        uri,
    ];

    let out = Command::new("youtube-dl")
        .args(&args)
        .stdin(Stdio::null())
        .output()?;

    if !out.status.success() {
        return Err(Error::Voice(VoiceError::YouTubeDLRun(out)));
    }

    let value = serde_json::from_reader(&out.stdout[..])?;
    let mut obj = match value {
        Value::Object(obj) => obj,
        other => return Err(Error::Voice(VoiceError::YouTubeDLProcessing(other))),
    };

    let uri = match obj.remove("url") {
        Some(v) => match v {
            Value::String(uri) => uri,
            other => return Err(Error::Voice(VoiceError::YouTubeDLUrl(other))),
        },
        None => return Err(Error::Voice(VoiceError::YouTubeDLUrl(Value::Object(obj)))),
    };

    ffmpeg(&uri)
}

fn is_stereo(path: &OsStr) -> Result<bool> {
    let args = ["-v", "quiet", "-of", "json", "-show-streams", "-i"];

    let out = Command::new("ffprobe")
        .args(&args)
        .arg(path)
        .stdin(Stdio::null())
        .output()?;

    let value: Value = serde_json::from_reader(&out.stdout[..])?;

    let streams = value
        .as_object()
        .and_then(|m| m.get("streams"))
        .and_then(|v| v.as_array())
        .ok_or(Error::Voice(VoiceError::Streams))?;

    let check = streams.iter().any(|stream| {
        let channels = stream
            .as_object()
            .and_then(|m| m.get("channels").and_then(|v| v.as_i64()));

        channels == Some(2)
    });

    Ok(check)
}
<commit_msg>Fewer ffprobe calls when playing audio through ffmpeg. (#472)<commit_after>use byteorder::{LittleEndian, ReadBytesExt};
use crate::internal::prelude::*;
use opus::{
    Channels,
    Decoder as OpusDecoder,
    Result as OpusResult,
};
use parking_lot::Mutex;
use serde_json;
use std::{
    ffi::OsStr,
    fs::File,
    io::{BufReader, ErrorKind as IoErrorKind, Read, Result as IoResult},
    process::{Child, Command, Stdio},
    result::Result as StdResult,
    sync::Arc,
};
use super::{AudioSource, AudioType, DcaError, DcaMetadata, VoiceError, audio};
use log::{debug, warn};

struct ChildContainer(Child);

impl Read for ChildContainer {
    fn read(&mut self, buffer: &mut [u8]) -> IoResult<usize> {
        self.0.stdout.as_mut().unwrap().read(buffer)
    }
}

impl Drop for ChildContainer {
    fn drop (&mut self) {
        if let Err(e) = self.0.kill() {
            debug!("[Voice] Error awaiting child process: {:?}", e);
        }
    }
}

// Since each audio item needs its own decoder, we need to
// work around the fact that OpusDecoders aint sendable.
struct SendDecoder(OpusDecoder);

impl SendDecoder {
    fn decode_float(&mut self, input: &[u8], output: &mut [f32], fec: bool) -> OpusResult<usize> {
        let &mut SendDecoder(ref mut sd) = self;
        sd.decode_float(input, output, fec)
    }
}

unsafe impl Send for SendDecoder {}

struct InputSource<R: Read + Send + 'static> {
    stereo: bool,
    reader: R,
    kind: AudioType,
    decoder: Option<Arc<Mutex<SendDecoder>>>,
}

impl<R: Read + Send> AudioSource for InputSource<R> {
    fn is_stereo(&mut self) -> bool { self.stereo }

    fn get_type(&self) -> AudioType { self.kind }

    fn read_pcm_frame(&mut self, buffer: &mut [i16]) -> Option<usize> {
        for (i, v) in buffer.iter_mut().enumerate() {
            *v = match self.reader.read_i16::<LittleEndian>() {
                Ok(v) => v,
                Err(ref e) => {
                    return if e.kind() == IoErrorKind::UnexpectedEof {
                        Some(i)
                    } else {
                        None
                    }
                },
            }
        }

        Some(buffer.len())
    }

    fn read_opus_frame(&mut self) -> Option<Vec<u8>> {
        match self.reader.read_i16::<LittleEndian>() {
            Ok(size) => {
                if size <= 0 {
                    warn!("Invalid opus frame size: {}", size);
                    return None;
                }

                let mut frame = Vec::with_capacity(size as usize);

                {
                    let reader = self.reader.by_ref();

                    if reader.take(size as u64).read_to_end(&mut frame).is_err() {
                        return None;
                    }
                }

                Some(frame)
            },
            Err(ref e) => if e.kind() == IoErrorKind::UnexpectedEof {
                Some(Vec::new())
            } else {
                None
            },
        }
    }

    fn decode_and_add_opus_frame(&mut self, float_buffer: &mut [f32; 1920], volume: f32) -> Option<usize> {
        let decoder_lock = self.decoder.as_mut()?.clone();
        let frame = self.read_opus_frame()?;
        let mut local_buf = [0f32; 960 * 2];

        let count = {
            let mut decoder = decoder_lock.lock();

            decoder.decode_float(frame.as_slice(), &mut local_buf, false).ok()?
        };

        for i in 0..1920 {
            float_buffer[i] += local_buf[i] * volume;
        }

        Some(count)
    }
}

/// Opens an audio file through `ffmpeg` and creates an audio source.
pub fn ffmpeg<P: AsRef<OsStr>>(path: P) -> Result<Box<dyn AudioSource>> {
    _ffmpeg(path.as_ref())
}

fn _ffmpeg(path: &OsStr) -> Result<Box<dyn AudioSource>> {
    // Will fail if the path is not to a file on the fs. Likely a YouTube URI.
    let is_stereo = is_stereo(path).unwrap_or(false);
    let stereo_val = if is_stereo { "2" } else { "1" };

    _ffmpeg_optioned(path, &[
        "-f",
        "s16le",
        "-ac",
        stereo_val,
        "-ar",
        "48000",
        "-acodec",
        "pcm_s16le",
        "-",
    ], Some(is_stereo))
}

/// Opens an audio file through `ffmpeg` and creates an audio source, with
/// user-specified arguments to pass to ffmpeg.
///
/// Note that this does _not_ build on the arguments passed by the [`ffmpeg`]
/// function.
///
/// # Examples
///
/// Pass options to create a custom ffmpeg streamer:
///
/// ```rust,no_run
/// use serenity::voice;
///
/// let stereo_val = "2";
///
/// let streamer = voice::ffmpeg_optioned("./some_file.mp3", &[
///     "-f",
///     "s16le",
///     "-ac",
///     stereo_val,
///     "-ar",
///     "48000",
///     "-acodec",
///     "pcm_s16le",
///     "-",
/// ]);
pub fn ffmpeg_optioned<P: AsRef<OsStr>>(
    path: P,
    args: &[&str],
) -> Result<Box<dyn AudioSource>> {
    _ffmpeg_optioned(path.as_ref(), args, None)
}

fn _ffmpeg_optioned(path: &OsStr, args: &[&str], is_stereo_known: Option<bool>) -> Result<Box<dyn AudioSource>> {
    let is_stereo = is_stereo_known
        .or_else(|| is_stereo(path).ok())
        .unwrap_or(false);

    let command = Command::new("ffmpeg")
        .arg("-i")
        .arg(path)
        .args(args)
        .stderr(Stdio::null())
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .spawn()?;

    Ok(pcm(is_stereo, ChildContainer(command)))
}

/// Creates a streamed audio source from a DCA file.
/// Currently only accepts the DCA1 format.
pub fn dca<P: AsRef<OsStr>>(path: P) -> StdResult<Box<dyn AudioSource>, DcaError> {
    _dca(path.as_ref())
}

fn _dca(path: &OsStr) -> StdResult<Box<dyn AudioSource>, DcaError> {
    let file = File::open(path).map_err(DcaError::IoError)?;

    let mut reader = BufReader::new(file);

    let mut header = [0u8; 4];

    // Read in the magic number to verify it's a DCA file.
    reader.read_exact(&mut header).map_err(DcaError::IoError)?;

    if header != b"DCA1"[..] {
        return Err(DcaError::InvalidHeader);
    }

    reader.read_exact(&mut header).map_err(DcaError::IoError)?;

    let size = (&header[..]).read_i32::<LittleEndian>().unwrap();

    // Sanity check
    if size < 2 {
        return Err(DcaError::InvalidSize(size));
    }

    let mut raw_json = Vec::with_capacity(size as usize);

    {
        let json_reader = reader.by_ref();
        json_reader
            .take(size as u64)
            .read_to_end(&mut raw_json)
            .map_err(DcaError::IoError)?;
    }

    let metadata = serde_json::from_slice::<DcaMetadata>(raw_json.as_slice())
        .map_err(DcaError::InvalidMetadata)?;

    Ok(opus(metadata.is_stereo(), reader))
}

/// Creates an Opus audio source. This makes certain assumptions: namely, that the input stream
/// is composed ONLY of opus frames of the variety that Discord expects.
///
/// If you want to decode a `.opus` file, use [`ffmpeg`]
///
/// [`ffmpeg`]: fn.ffmpeg.html
pub fn opus<R: Read + Send + 'static>(is_stereo: bool, reader: R) -> Box<dyn AudioSource> {
    Box::new(InputSource {
        stereo: is_stereo,
        reader,
        kind: AudioType::Opus,
        decoder: Some(
            Arc::new(Mutex::new(
                // We always want to decode *to* stereo, for mixing reasons.
                SendDecoder(OpusDecoder::new(audio::SAMPLE_RATE, Channels::Stereo).unwrap())
            ))
        ),
    })
}

/// Creates a PCM audio source.
pub fn pcm<R: Read + Send + 'static>(is_stereo: bool, reader: R) -> Box<dyn AudioSource> {
    Box::new(InputSource {
        stereo: is_stereo,
        reader,
        kind: AudioType::Pcm,
        decoder: None,
    })
}

/// Creates a streamed audio source with `youtube-dl` and `ffmpeg`.
pub fn ytdl(uri: &str) -> Result<Box<dyn AudioSource>> {
    let args = [
        "-f",
        "webm[abr>0]/bestaudio/best",
        "--no-playlist",
        "--print-json",
        "--skip-download",
        uri,
    ];

    let out = Command::new("youtube-dl")
        .args(&args)
        .stdin(Stdio::null())
        .output()?;

    if !out.status.success() {
        return Err(Error::Voice(VoiceError::YouTubeDLRun(out)));
    }

    let value = serde_json::from_reader(&out.stdout[..])?;
    let mut obj = match value {
        Value::Object(obj) => obj,
        other => return Err(Error::Voice(VoiceError::YouTubeDLProcessing(other))),
    };

    let uri = match obj.remove("url") {
        Some(v) => match v {
            Value::String(uri) => uri,
            other => return Err(Error::Voice(VoiceError::YouTubeDLUrl(other))),
        },
        None => return Err(Error::Voice(VoiceError::YouTubeDLUrl(Value::Object(obj)))),
    };

    ffmpeg(&uri)
}

fn is_stereo(path: &OsStr) -> Result<bool> {
    let args = ["-v", "quiet", "-of", "json", "-show-streams", "-i"];

    let out = Command::new("ffprobe")
        .args(&args)
        .arg(path)
        .stdin(Stdio::null())
        .output()?;

    let value: Value = serde_json::from_reader(&out.stdout[..])?;

    let streams = value
        .as_object()
        .and_then(|m| m.get("streams"))
        .and_then(|v| v.as_array())
        .ok_or(Error::Voice(VoiceError::Streams))?;

    let check = streams.iter().any(|stream| {
        let channels = stream
            .as_object()
            .and_then(|m| m.get("channels").and_then(|v| v.as_i64()));

        channels == Some(2)
    });

    Ok(check)
}
<|endoftext|>


<commit_before>
//   Copyright 2014 Colin Sherratt
//
//   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.

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

use cow::join::{join_set_to_map};

use libc::c_void;

use gl;
use gl::types::{GLsizei, GLuint};

use cgmath::Matrix4;
use cgmath::Array2;
use cgmath::{Vector, EuclideanVector};

use config::Config;
use graphics::Graphics;
use snowmew::common::Entity;

use db::GlState;

#[repr(packed)]
struct DrawElementsIndirectCommand {
    pub count: GLuint,
    pub instrance_count: GLuint,
    pub first_index: GLuint,
    pub base_vertex: GLuint,
    pub base_instance: GLuint
}

pub struct CommandBufferIndirect {
    command: GLuint,
    ptr: *mut DrawElementsIndirectCommand,
    size: uint,
    batches: Vec<Batch>
}

#[deriving(Clone)]
pub struct Batch {
    vbo: Entity,
    offset: uint,
    count: uint
}

impl Batch {
    pub fn vbo(&self) -> Entity {self.vbo}

    pub fn offset(&self) -> *const c_void {
        assert!(mem::size_of::<DrawElementsIndirectCommand>() == 20);
        (self.offset * mem::size_of::<DrawElementsIndirectCommand>()) as *const c_void
    }

    pub fn drawcount(&self) -> GLsizei {
        self.count as GLsizei
    }

    pub fn stride(&self) -> GLsizei {
        assert!(mem::size_of::<DrawElementsIndirectCommand>() == 20);
        mem::size_of::<DrawElementsIndirectCommand>() as GLsizei
    }

    pub fn offset_int(&self) -> uint {self.offset}
}

impl CommandBufferIndirect {
    pub fn new(cfg: &Config) -> CommandBufferIndirect {
        let cb = &mut [0];
        unsafe {
            gl::GenBuffers(1, cb.unsafe_mut(0));
            gl::BindBuffer(gl::DRAW_INDIRECT_BUFFER, cb[0]);
            gl::BufferData(gl::DRAW_INDIRECT_BUFFER,
                           (mem::size_of::<DrawElementsIndirectCommand>() *
                           cfg.max_size()) as i64,
                           ptr::null(),
                           gl::DYNAMIC_DRAW);
            gl::BindBuffer(gl::DRAW_INDIRECT_BUFFER, 0);
        }

        CommandBufferIndirect {
            command: cb[0],
            ptr: ptr::null_mut(),
            size: cfg.max_size(),
            batches: Vec::new()
        }
    }

    pub fn map(&mut self) { unsafe {
        gl::BindBuffer(gl::DRAW_INDIRECT_BUFFER, self.command);
        self.ptr = gl::MapBufferRange(
            gl::DRAW_INDIRECT_BUFFER, 0,
            (mem::size_of::<DrawElementsIndirectCommand>() *
            self.size) as i64,
            gl::MAP_WRITE_BIT | gl::MAP_INVALIDATE_BUFFER_BIT
        ) as *mut DrawElementsIndirectCommand;
        gl::BindBuffer(gl::DRAW_INDIRECT_BUFFER, 0);
    }}

    pub fn unmap(&mut self) { unsafe {
        self.ptr = ptr::null_mut();
        gl::BindBuffer(gl::DRAW_INDIRECT_BUFFER, self.command);
        gl::UnmapBuffer(gl::DRAW_INDIRECT_BUFFER);
        gl::BindBuffer(gl::DRAW_INDIRECT_BUFFER, 0);
    }}

    pub fn build<GD: Graphics>(&mut self, db: &GD, scene: Entity, instanced_is_enabled: bool) {
        let mut batch = Batch {
            vbo: 0,
            offset: 0,
            count: 0
        };

        let mut idx = -1;
        let mut last_geo = None;
        let mut command = DrawElementsIndirectCommand {
            count: 0,
            instrance_count: 1,
            first_index: 0,
            base_vertex: 0,
            base_instance: 0
        };

        unsafe {
            self.batches.truncate(0);
            let b = slice::from_raw_mut_buf(&self.ptr, self.size);
            for (count, (_, draw)) in join_set_to_map(db.scene_iter(scene), db.drawable_iter()).enumerate() {
                if idx == -1 {
                    let draw_geo = db.geometry(draw.geometry).expect("geometry not found");
                    last_geo = Some(draw.geometry);
                    command = DrawElementsIndirectCommand {
                        count: draw_geo.count as GLuint,
                        instrance_count: 1,
                        first_index: draw_geo.offset as GLuint,
                        base_vertex: 0,
                        base_instance: count as GLuint
                    };

                    batch.vbo = draw_geo.vb;
                    batch.count = 1;

                    idx = 0;
                } else if last_geo == Some(draw.geometry) && instanced_is_enabled {
                    command.instrance_count += 1;
                } else {
                    let draw_geo = db.geometry(draw.geometry).expect("geometry not found");
                    last_geo = Some(draw.geometry);

                    b[idx] = command;
                    idx += 1;

                    command = DrawElementsIndirectCommand {
                        count: draw_geo.count as GLuint,
                        instrance_count: 1,
                        first_index: draw_geo.offset as GLuint,
                        base_vertex: 0,
                        base_instance: count as GLuint
                    };

                    if batch.vbo == 0 {
                        batch.vbo = draw_geo.vb;
                        batch.offset = idx;
                        batch.count = 1;
                    } else if batch.vbo == draw_geo.vb {
                        batch.count += 1;
                    } else {
                        self.batches.push(batch.clone());
                        batch.vbo = draw_geo.vb;
                        batch.offset = idx;
                        batch.count = 1;
                    }

                }
            }
            b[idx] = command;
        }
        self.batches.push(batch)
    }

    pub fn cull(&self, draw: GLuint, matrix: GLuint, dat: &GlState, mat: &Matrix4<f32>) {
        let to_plane = |x, scale| {
            let plane = mat.row(x).mul_s(scale).add_v(&mat.row(3));
            plane.normalize()
        };

        let planes = &[
            to_plane(0,  1.),
            to_plane(0, -1.),
            to_plane(1,  1.),
            to_plane(1, -1.),
            to_plane(2,  1.),
            to_plane(2, -1.)
        ];

        let shader = dat.compute_cull.as_ref().expect("Could not get cull");

        let size = self.batches.iter().fold(0, |a, b| a + b.count);

        let x = 256i;
        let y = size / 256 + 1;

        shader.bind();
        unsafe {
            gl::Uniform4fv(shader.uniform("plane"), 6, &planes[0].x);
            gl::Uniform1i(shader.uniform("max_id"), size as i32);
            gl::BindBufferBase(gl::SHADER_STORAGE_BUFFER, 0, draw);
            gl::BindBufferBase(gl::SHADER_STORAGE_BUFFER, 1, matrix);
            gl::BindBufferBase(gl::SHADER_STORAGE_BUFFER, 2, self.command);
            gl::DispatchCompute(x as u32, y as u32, 1);
            gl::MemoryBarrier(gl::COMMAND_BARRIER_BIT);
        }
    }

    pub fn batches<'a>(&'a self) -> &'a [Batch] {
        self.batches.slice(0, self.batches.len())
    }

    pub fn id(&self) -> GLuint {
        self.command
    }
}

pub struct CommandBufferEmulated {
    commands: Vec<DrawElementsIndirectCommand>,
    batches: Vec<Batch>
}

impl CommandBufferEmulated {
    pub fn new(_: &Config) -> CommandBufferEmulated {
        CommandBufferEmulated {
            commands: Vec::new(),
            batches: Vec::new()
        }
    }

    pub fn map(&mut self) {}
    pub fn unmap(&mut self) {}

    pub fn build<GD: Graphics>(&mut self, db: &GD, scene: Entity, _: bool) {
        let mut batch = Batch {
            vbo: 0,
            offset: 0,
            count: 0
        };

        self.batches.truncate(0);
        self.commands.truncate(0);
        for (count, (_, draw)) in join_set_to_map(db.scene_iter(scene), db.drawable_iter()).enumerate() {
            let draw_geo = db.geometry(draw.geometry).expect("geometry not found");

            self.commands.push(DrawElementsIndirectCommand {
                count: draw_geo.count as GLuint,
                instrance_count: 1,
                first_index: draw_geo.offset as GLuint,
                base_vertex: 0,
                base_instance: count as GLuint
            });

            if batch.vbo == 0 {
                batch.vbo = draw_geo.vb;
                batch.offset = count;
                batch.count = 1;
            } else if batch.vbo == draw_geo.vb {
                batch.count += 1;
            } else {
                self.batches.push(batch.clone());
                batch.vbo = draw_geo.vb;
                batch.offset = count;
                batch.count = 1;
            }
        }

        self.batches.push(batch)
    }

    pub fn batches<'a>(&'a self) -> &'a [Batch] {
        self.batches.slice(0, self.batches.len())
    }

    pub fn commands<'a>(&'a self) -> &'a [DrawElementsIndirectCommand] {
        self.commands.slice(0, self.commands.len())
    }
}<commit_msg>fix instancing for fallback command queue<commit_after>
//   Copyright 2014 Colin Sherratt
//
//   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.

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

use cow::join::{join_set_to_map};

use libc::c_void;

use gl;
use gl::types::{GLsizei, GLuint};

use cgmath::Matrix4;
use cgmath::Array2;
use cgmath::{Vector, EuclideanVector};

use config::Config;
use graphics::Graphics;
use snowmew::common::Entity;

use db::GlState;

#[repr(packed)]
struct DrawElementsIndirectCommand {
    pub count: GLuint,
    pub instrance_count: GLuint,
    pub first_index: GLuint,
    pub base_vertex: GLuint,
    pub base_instance: GLuint
}

pub struct CommandBufferIndirect {
    command: GLuint,
    ptr: *mut DrawElementsIndirectCommand,
    size: uint,
    batches: Vec<Batch>
}

#[deriving(Clone)]
pub struct Batch {
    vbo: Entity,
    offset: uint,
    count: uint
}

impl Batch {
    pub fn vbo(&self) -> Entity {self.vbo}

    pub fn offset(&self) -> *const c_void {
        assert!(mem::size_of::<DrawElementsIndirectCommand>() == 20);
        (self.offset * mem::size_of::<DrawElementsIndirectCommand>()) as *const c_void
    }

    pub fn drawcount(&self) -> GLsizei {
        self.count as GLsizei
    }

    pub fn stride(&self) -> GLsizei {
        assert!(mem::size_of::<DrawElementsIndirectCommand>() == 20);
        mem::size_of::<DrawElementsIndirectCommand>() as GLsizei
    }

    pub fn offset_int(&self) -> uint {self.offset}
}

impl CommandBufferIndirect {
    pub fn new(cfg: &Config) -> CommandBufferIndirect {
        let cb = &mut [0];
        unsafe {
            gl::GenBuffers(1, cb.unsafe_mut(0));
            gl::BindBuffer(gl::DRAW_INDIRECT_BUFFER, cb[0]);
            gl::BufferData(gl::DRAW_INDIRECT_BUFFER,
                           (mem::size_of::<DrawElementsIndirectCommand>() *
                           cfg.max_size()) as i64,
                           ptr::null(),
                           gl::DYNAMIC_DRAW);
            gl::BindBuffer(gl::DRAW_INDIRECT_BUFFER, 0);
        }

        CommandBufferIndirect {
            command: cb[0],
            ptr: ptr::null_mut(),
            size: cfg.max_size(),
            batches: Vec::new()
        }
    }

    pub fn map(&mut self) { unsafe {
        gl::BindBuffer(gl::DRAW_INDIRECT_BUFFER, self.command);
        self.ptr = gl::MapBufferRange(
            gl::DRAW_INDIRECT_BUFFER, 0,
            (mem::size_of::<DrawElementsIndirectCommand>() *
            self.size) as i64,
            gl::MAP_WRITE_BIT | gl::MAP_INVALIDATE_BUFFER_BIT
        ) as *mut DrawElementsIndirectCommand;
        gl::BindBuffer(gl::DRAW_INDIRECT_BUFFER, 0);
    }}

    pub fn unmap(&mut self) { unsafe {
        self.ptr = ptr::null_mut();
        gl::BindBuffer(gl::DRAW_INDIRECT_BUFFER, self.command);
        gl::UnmapBuffer(gl::DRAW_INDIRECT_BUFFER);
        gl::BindBuffer(gl::DRAW_INDIRECT_BUFFER, 0);
    }}

    pub fn build<GD: Graphics>(&mut self, db: &GD, scene: Entity, instanced_is_enabled: bool) {
        let mut batch = Batch {
            vbo: 0,
            offset: 0,
            count: 0
        };

        let mut idx = -1;
        let mut last_geo = None;
        let mut command = DrawElementsIndirectCommand {
            count: 0,
            instrance_count: 1,
            first_index: 0,
            base_vertex: 0,
            base_instance: 0
        };

        unsafe {
            self.batches.truncate(0);
            let b = slice::from_raw_mut_buf(&self.ptr, self.size);
            for (count, (_, draw)) in join_set_to_map(db.scene_iter(scene), db.drawable_iter()).enumerate() {
                if idx == -1 {
                    let draw_geo = db.geometry(draw.geometry).expect("geometry not found");
                    last_geo = Some(draw.geometry);
                    command = DrawElementsIndirectCommand {
                        count: draw_geo.count as GLuint,
                        instrance_count: 1,
                        first_index: draw_geo.offset as GLuint,
                        base_vertex: 0,
                        base_instance: count as GLuint
                    };

                    batch.vbo = draw_geo.vb;
                    batch.count = 1;

                    idx = 0;
                } else if last_geo == Some(draw.geometry) && instanced_is_enabled {
                    command.instrance_count += 1;
                } else {
                    let draw_geo = db.geometry(draw.geometry).expect("geometry not found");
                    last_geo = Some(draw.geometry);

                    b[idx] = command;
                    idx += 1;

                    command = DrawElementsIndirectCommand {
                        count: draw_geo.count as GLuint,
                        instrance_count: 1,
                        first_index: draw_geo.offset as GLuint,
                        base_vertex: 0,
                        base_instance: count as GLuint
                    };

                    if batch.vbo == 0 {
                        batch.vbo = draw_geo.vb;
                        batch.offset = idx;
                        batch.count = 1;
                    } else if batch.vbo == draw_geo.vb {
                        batch.count += 1;
                    } else {
                        self.batches.push(batch.clone());
                        batch.vbo = draw_geo.vb;
                        batch.offset = idx;
                        batch.count = 1;
                    }

                }
            }
            b[idx] = command;
        }
        self.batches.push(batch)
    }

    pub fn cull(&self, draw: GLuint, matrix: GLuint, dat: &GlState, mat: &Matrix4<f32>) {
        let to_plane = |x, scale| {
            let plane = mat.row(x).mul_s(scale).add_v(&mat.row(3));
            plane.normalize()
        };

        let planes = &[
            to_plane(0,  1.),
            to_plane(0, -1.),
            to_plane(1,  1.),
            to_plane(1, -1.),
            to_plane(2,  1.),
            to_plane(2, -1.)
        ];

        let shader = dat.compute_cull.as_ref().expect("Could not get cull");

        let size = self.batches.iter().fold(0, |a, b| a + b.count);

        let x = 256i;
        let y = size / 256 + 1;

        shader.bind();
        unsafe {
            gl::Uniform4fv(shader.uniform("plane"), 6, &planes[0].x);
            gl::Uniform1i(shader.uniform("max_id"), size as i32);
            gl::BindBufferBase(gl::SHADER_STORAGE_BUFFER, 0, draw);
            gl::BindBufferBase(gl::SHADER_STORAGE_BUFFER, 1, matrix);
            gl::BindBufferBase(gl::SHADER_STORAGE_BUFFER, 2, self.command);
            gl::DispatchCompute(x as u32, y as u32, 1);
            gl::MemoryBarrier(gl::COMMAND_BARRIER_BIT);
        }
    }

    pub fn batches<'a>(&'a self) -> &'a [Batch] {
        self.batches.slice(0, self.batches.len())
    }

    pub fn id(&self) -> GLuint {
        self.command
    }
}

pub struct CommandBufferEmulated {
    commands: Vec<DrawElementsIndirectCommand>,
    batches: Vec<Batch>
}

impl CommandBufferEmulated {
    pub fn new(_: &Config) -> CommandBufferEmulated {
        CommandBufferEmulated {
            commands: Vec::new(),
            batches: Vec::new()
        }
    }

    pub fn map(&mut self) {}
    pub fn unmap(&mut self) {}

    pub fn build<GD: Graphics>(&mut self, db: &GD, scene: Entity, _: bool) {
        let mut batch = Batch {
            vbo: 0,
            offset: 0,
            count: 0
        };

        self.batches.truncate(0);
        self.commands.truncate(0);
        for (count, (_, draw)) in join_set_to_map(db.scene_iter(scene), db.drawable_iter()).enumerate() {
            let draw_geo = db.geometry(draw.geometry).expect("geometry not found");

            let len = self.commands.len();
            let cmd = if batch.vbo == draw_geo.vb && len != 0 {
                let last = &mut self.commands[len-1];
                if last.count == draw_geo.count as u32 &&
                   last.first_index == draw_geo.offset as u32 {
                    last.instrance_count += 1;
                    None
                } else {
                    Some(DrawElementsIndirectCommand {
                        count: draw_geo.count as GLuint,
                        instrance_count: 1,
                        first_index: draw_geo.offset as GLuint,
                        base_vertex: 0,
                        base_instance: count as GLuint
                    })
                }
            } else {
                Some(DrawElementsIndirectCommand {
                    count: draw_geo.count as GLuint,
                    instrance_count: 1,
                    first_index: draw_geo.offset as GLuint,
                    base_vertex: 0,
                    base_instance: count as GLuint
                })
            };

            if let Some(cmd) = cmd { self.commands.push(cmd); }

            if batch.vbo == 0 {
                batch.vbo = draw_geo.vb;
                batch.offset = count;
                batch.count = 1;
            } else if batch.vbo != draw_geo.vb {
                self.batches.push(batch.clone());
                batch.vbo = draw_geo.vb;
                batch.offset = count;
                batch.count = 1;
            }
        }

        self.batches.push(batch)
    }

    pub fn batches<'a>(&'a self) -> &'a [Batch] {
        self.batches.slice(0, self.batches.len())
    }

    pub fn commands<'a>(&'a self) -> &'a [DrawElementsIndirectCommand] {
        self.commands.slice(0, self.commands.len())
    }
}<|endoftext|>


<commit_before>
struct median_filter_t<'a> {
    window: &'a mut u8,
    input: &'a mut usize,
    hole: &'a mut i64,
    max: &'a mut i64,
}

fn median_init(f: median_filter_t, w_size: isize){
    *f.max = 0;
    *f.hole = 0;
}

fn main() {
    println!("Hello, world!");
}
<commit_msg>remove useless parameter<commit_after>
struct median_filter_t<'a> {
    window: &'a mut u8,
    input: &'a mut usize,
    hole: &'a mut i64,
    max: &'a mut i64,
}

fn median_init(f: median_filter_t){
    *f.max = 0;
    *f.hole = 0;
}

fn main() {
    println!("Hello, world!");
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate bitflags;

#[macro_use]
extern crate lazy_static;

extern crate winapi;
pub extern crate stretch;

#[cfg(feature="all")]
#[cfg(test)]
mod tests;

mod errors;
pub use errors::{NwgError};

mod events;
pub use events::*;

mod common_types;
pub use common_types::*;

pub(crate) mod win32;
pub use win32::{
 dispatch_thread_events, dispatch_thread_events_with_callback, stop_thread_dispatch, enable_visual_styles, init_common_controls,
 window::{
     EventHandler, RawEventHandler,
     full_bind_event_handler, bind_event_handler, unbind_event_handler,
     bind_raw_event_handler, has_raw_handler, unbind_raw_event_handler
 },
 message_box::{MessageButtons, MessageIcons, MessageChoice, MessageParams, message, fatal_message, error_message, simple_message}};

pub(crate) use win32::window::bind_raw_event_handler_inner;

#[allow(deprecated)]
pub use win32::high_dpi::{set_dpi_awareness, scale_factor, dpi};

#[cfg(feature="cursor")]
pub use win32::cursor::GlobalCursor;

#[cfg(feature="clipboard")]
pub use win32::clipboard::{Clipboard, ClipboardFormat, ClipboardData};

mod resources;
pub use resources::*;

mod controls;
pub use controls::*;

mod layouts;
pub use layouts::*;

#[cfg(feature = "winnls")]
mod winnls;

#[cfg(feature = "winnls")]
pub use winnls::*;

/**
    A structure that implements this trait is considered a GUI structure. The structure will hold GUI components and possibly user data.

    A structure that implements `PartialUi` must be part of another UI structure and cannot be used at it is. It will most likely used
    as a struct member of another struct that implements `NativeUi`.

    The goal of `NativeUi` and `PartialUi` is to provide a common way to define NWG applications.
    Native-windows-derive can automatically implements this trait.

    For an example on how to implement this trait, see the **Small application layout** section in the NWG documentation.
*/
pub trait PartialUi {
    /**
        Should initializes the GUI components. Similar to `NativeUi::build_ui` except it doesn't handle events binding.

        Parameters:
          - `data`: A reference to the struct data from the parent struct
          - `parent`: An optional reference to the parent UI control. If this is defined, ui controls of the partial should be children of this value.
    */
    fn build_partial<W: Into<ControlHandle>>(data: &mut Self, parent: Option<W>) -> Result<(), NwgError>;

    /**
        Should process the events of the partial. This method will probably be called from an event handler bound in the parent GUI structure.

        Parameters:
          - `base`: A reference to the parent struct data
          - `evt`: The event raised
          - `evt_data`: The data of the event raised
          - `handle`: Handle of the control that raised the event
    */
    fn process_event(&self, _evt: Event, _evt_data: &EventData, _handle: ControlHandle) {}

    /**
        Should return the handles of the top level parent controls (such as Windows). Those handle should be used to bind
        the default events handler.
    */
    fn handles<'a>(&'a self) -> Vec<&'a ControlHandle> { vec![] }
}

/**
    A structure that implements this trait is considered a GUI structure. The structure will hold GUI components and possibly user data.

    The goal of `NativeUi` and `PartialUi` is to provide a common way to define NWG applications.
    Native-windows-derive can automatically implements this trait.

    For an example on how to implement this trait, see the **Small application layout** section in the NWG documentation.
*/
pub trait NativeUi<UI> {

    /**
        A constructor for the structure. It should initializes the GUI components and bind GUI events.

        Parameters:
          - `inital_state`: should contain the initial user data. `NativeUi` assumes this data will be wrapped in another type (`UI`).
    */
    fn build_ui(inital_state: Self) -> Result<UI, NwgError>;
}

/// Initialize some application wide GUI settings.
/// This includes default styling and common controls resources.
pub fn init() -> std::result::Result<(), errors::NwgError> {
    if cfg!(not(feature="no-styling")) {
        enable_visual_styles();
    }

    init_common_controls()
}
<commit_msg>smallfix<commit_after>#[macro_use]
extern crate bitflags;

#[macro_use]
extern crate lazy_static;

extern crate winapi;

#[cfg(feature="flexbox")]
pub extern crate stretch;

#[cfg(feature="all")]
#[cfg(test)]
mod tests;

mod errors;
pub use errors::{NwgError};

mod events;
pub use events::*;

mod common_types;
pub use common_types::*;

pub(crate) mod win32;
pub use win32::{
 dispatch_thread_events, dispatch_thread_events_with_callback, stop_thread_dispatch, enable_visual_styles, init_common_controls,
 window::{
     EventHandler, RawEventHandler,
     full_bind_event_handler, bind_event_handler, unbind_event_handler,
     bind_raw_event_handler, has_raw_handler, unbind_raw_event_handler
 },
 message_box::{MessageButtons, MessageIcons, MessageChoice, MessageParams, message, fatal_message, error_message, simple_message}};

pub(crate) use win32::window::bind_raw_event_handler_inner;

#[allow(deprecated)]
pub use win32::high_dpi::{set_dpi_awareness, scale_factor, dpi};

#[cfg(feature="cursor")]
pub use win32::cursor::GlobalCursor;

#[cfg(feature="clipboard")]
pub use win32::clipboard::{Clipboard, ClipboardFormat, ClipboardData};

mod resources;
pub use resources::*;

mod controls;
pub use controls::*;

mod layouts;
pub use layouts::*;

#[cfg(feature = "winnls")]
mod winnls;

#[cfg(feature = "winnls")]
pub use winnls::*;

/**
    A structure that implements this trait is considered a GUI structure. The structure will hold GUI components and possibly user data.

    A structure that implements `PartialUi` must be part of another UI structure and cannot be used at it is. It will most likely used
    as a struct member of another struct that implements `NativeUi`.

    The goal of `NativeUi` and `PartialUi` is to provide a common way to define NWG applications.
    Native-windows-derive can automatically implements this trait.

    For an example on how to implement this trait, see the **Small application layout** section in the NWG documentation.
*/
pub trait PartialUi {
    /**
        Should initializes the GUI components. Similar to `NativeUi::build_ui` except it doesn't handle events binding.

        Parameters:
          - `data`: A reference to the struct data from the parent struct
          - `parent`: An optional reference to the parent UI control. If this is defined, ui controls of the partial should be children of this value.
    */
    fn build_partial<W: Into<ControlHandle>>(data: &mut Self, parent: Option<W>) -> Result<(), NwgError>;

    /**
        Should process the events of the partial. This method will probably be called from an event handler bound in the parent GUI structure.

        Parameters:
          - `base`: A reference to the parent struct data
          - `evt`: The event raised
          - `evt_data`: The data of the event raised
          - `handle`: Handle of the control that raised the event
    */
    fn process_event(&self, _evt: Event, _evt_data: &EventData, _handle: ControlHandle) {}

    /**
        Should return the handles of the top level parent controls (such as Windows). Those handle should be used to bind
        the default events handler.
    */
    fn handles<'a>(&'a self) -> Vec<&'a ControlHandle> { vec![] }
}

/**
    A structure that implements this trait is considered a GUI structure. The structure will hold GUI components and possibly user data.

    The goal of `NativeUi` and `PartialUi` is to provide a common way to define NWG applications.
    Native-windows-derive can automatically implements this trait.

    For an example on how to implement this trait, see the **Small application layout** section in the NWG documentation.
*/
pub trait NativeUi<UI> {

    /**
        A constructor for the structure. It should initializes the GUI components and bind GUI events.

        Parameters:
          - `inital_state`: should contain the initial user data. `NativeUi` assumes this data will be wrapped in another type (`UI`).
    */
    fn build_ui(inital_state: Self) -> Result<UI, NwgError>;
}

/// Initialize some application wide GUI settings.
/// This includes default styling and common controls resources.
pub fn init() -> std::result::Result<(), errors::NwgError> {
    if cfg!(not(feature="no-styling")) {
        enable_visual_styles();
    }

    init_common_controls()
}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io
use super::query;
use hyper::mime::Mime;
use iron::status::Status as IronStatus;
use iron::typemap::Key;
use mimir::rubber::Rubber;
use model;
use model::v1::*;
use params::{
    coord_param, dataset_param, get_param_array, paginate_param, shape_param, types_param,
};
use prometheus;
use prometheus::Encoder;
use rustless;
use rustless::server::header;
use rustless::{Api, Nesting};
use serde;
use serde_json;
use valico::json_dsl;

use navitia_model::objects::Coord;

const DEFAULT_LIMIT: u64 = 10u64;
const DEFAULT_OFFSET: u64 = 0u64;

struct Timer;
impl Key for Timer {
    type Value = prometheus::HistogramTimer;
}

lazy_static! {
    static ref HTTP_COUNTER: prometheus::CounterVec = register_counter_vec!(
        "bragi_http_requests_total",
        "Total number of HTTP requests made.",
        &["handler", "method", "status"]
    ).unwrap();
    static ref HTTP_REQ_HISTOGRAM: prometheus::HistogramVec = register_histogram_vec!(
        "bragi_http_request_duration_seconds",
        "The HTTP request latencies in seconds.",
        &["handler", "method"],
        prometheus::exponential_buckets(0.001, 1.5, 25).unwrap()
    ).unwrap();
}

fn render<T>(
    mut client: rustless::Client,
    obj: T,
) -> Result<rustless::Client, rustless::ErrorResponse>
where
    T: serde::Serialize + model::v1::HasStatus,
{
    client.set_json_content_type();
    client.set_header(header::AccessControlAllowOrigin::Any);
    client.set_status(obj.status());
    client.text(serde_json::to_string(&obj).unwrap())
}

pub struct ApiEndPoint {
    pub es_cnx_string: String,
}

impl ApiEndPoint {
    pub fn root(&self) -> rustless::Api {
        Api::build(|api| {
            api.get("", |endpoint| {
                endpoint.handle(|client, _params| {
                    let desc = EndPoint {
                        description: "autocomplete service".to_string(),
                    };
                    render(client, desc)
                })
            });

            api.error_formatter(|error, _media| {
                let err = if error.is::<rustless::errors::Validation>() {
                    let val_err = error.downcast::<rustless::errors::Validation>().unwrap();
                    // TODO better message, we shouldn't use {:?} but access the `path`
                    // and `detail` of all errrors in val_err.reason
                    CustomError {
                        short: "validation error".to_string(),
                        long: format!("invalid arguments {:?}", val_err.reason),
                        status: IronStatus::BadRequest,
                    }
                } else {
                    CustomError {
                        short: "bad_request".to_string(),
                        long: format!("bad request, error: {}", error),
                        status: IronStatus::BadRequest,
                    }
                };
                let mut resp = rustless::Response::from(
                    err.status,
                    Box::new(serde_json::to_string(&err).unwrap()),
                );
                resp.set_json_content_type();
                Some(resp)
            });

            api.before(|client, _params| {
                let method = client.endpoint.method.to_string();

                HTTP_REQ_HISTOGRAM
                    .get_metric_with(&labels!{
                        "handler" => client.endpoint.path.path.as_str(),
                        "method" => method.as_str(),
                    })
                    .map(|timer| {
                        client.ext.insert::<Timer>(timer.start_timer());
                    })
                    .unwrap_or_else(|err| {
                        error!("impossible to get HTTP_REQ_HISTOGRAM metrics";
                               "err" => err.to_string());
                    });
                Ok(())
            });

            api.after(|client, _params| {
                let method = client.endpoint.method.to_string();
                let code = client.status().to_string();

                HTTP_COUNTER
                    .get_metric_with(&labels!{
                        "handler" => client.endpoint.path.path.as_str(),
                        "method" => method.as_str(),
                        "status" => code.as_str(),
                    })
                    .map(|counter| counter.inc())
                    .unwrap_or_else(|err| {
                        error!("impossible to get HTTP_COUNTER metrics"; "err" => err.to_string());
                    });
                client
                    .ext
                    .remove::<Timer>()
                    .map(|timer| timer.observe_duration())
                    .unwrap_or_else(|| error!("impossible to get timers from typemap"));
                Ok(())
            });
            api.mount(self.v1());
        })
    }

    fn v1(&self) -> rustless::Api {
        Api::build(|api| {
            api.mount(self.status());
            api.mount(self.autocomplete());
            api.mount(self.features());
            api.mount(self.reverse());
            api.mount(self.metrics());
        })
    }

    fn status(&self) -> rustless::Api {
        Api::build(|api| {
            api.get("status", |endpoint| {
                let cnx = self.es_cnx_string.clone();
                endpoint.handle(move |client, _params| {
                    let status = Status {
                        version: env!("CARGO_PKG_VERSION").to_string(),
                        es: cnx.to_string(),
                        status: "good".to_string(),
                    };
                    render(client, status)
                })
            });
        })
    }

    fn metrics(&self) -> rustless::Api {
        Api::build(|api| {
            api.get("metrics", |endpoint| {
                endpoint.handle(move |mut client, _params| {
                    let encoder = prometheus::TextEncoder::new();
                    let metric_familys = prometheus::gather();
                    let mut buffer = vec![];
                    encoder.encode(&metric_familys, &mut buffer).unwrap();
                    client.set_content_type(encoder.format_type().parse::<Mime>().unwrap());
                    client.text(String::from_utf8(buffer).unwrap())
                })
            });
        })
    }

    fn reverse(&self) -> rustless::Api {
        Api::build(|api| {
            api.get("reverse", |endpoint| {
                endpoint.params(|params| {
                    coord_param(params, false);
                });
                let cnx = self.es_cnx_string.clone();
                endpoint.handle(move |client, params| {
                    let coord = ::mimir::Coord::new(
                        params.find("lon").and_then(|p| p.as_f64()).unwrap(),
                        params.find("lat").and_then(|p| p.as_f64()).unwrap(),
                    );
                    let mut rubber = Rubber::new(&cnx);
                    let model_autocomplete =
                        rubber.get_address(&coord).map_err(model::BragiError::from);

                    let response = model::v1::AutocompleteResponse::from(model_autocomplete);
                    render(client, response)
                })
            });
        })
    }

    fn features(&self) -> rustless::Api {
        Api::build(|api| {
            api.get("features/:id", |endpoint| {
                endpoint.params(|params| {
                    params.opt_typed("id", json_dsl::string());
                    dataset_param(params);
                });

                let cnx = self.es_cnx_string.clone();
                endpoint.handle(move |client, params| {
                    let id = params.find("id").unwrap().as_str().unwrap();
                    let pt_datasets = get_param_array(params, "pt_dataset");
                    let all_data = params
                        .find("_all_data")
                        .map_or(false, |val| val.as_bool().unwrap());
                    let features = query::features(&pt_datasets, all_data, &cnx, &id);
                    let response = model::v1::AutocompleteResponse::from(features);
                    render(client, response)
                })
            });
        })
    }

    fn autocomplete(&self) -> rustless::Api {
        Api::build(|api| {
            api.post("autocomplete", |endpoint| {
                endpoint.params(|params| {
                    params.opt_typed("q", json_dsl::string());
                    dataset_param(params);
                    paginate_param(params);
                    shape_param(params);
                    types_param(params);
                });

                let cnx = self.es_cnx_string.clone();
                endpoint.handle(move |client, params| {
                    let q = params
                        .find("q")
                        .and_then(|val| val.as_str())
                        .unwrap_or("")
                        .to_string();
                    let pt_datasets = get_param_array(params, "pt_dataset");
                    let all_data = params
                        .find("_all_data")
                        .map_or(false, |val| val.as_bool().unwrap());
                    let offset = params
                        .find("offset")
                        .and_then(|val| val.as_u64())
                        .unwrap_or(DEFAULT_OFFSET);
                    let limit = params
                        .find("limit")
                        .and_then(|val| val.as_u64())
                        .unwrap_or(DEFAULT_LIMIT);
                    let geometry = params.find_path(&["shape", "geometry"]).unwrap();
                    let coordinates = geometry
                        .find_path(&["coordinates"])
                        .unwrap()
                        .as_array()
                        .unwrap();
                    let mut shape = Vec::new();
                    for ar in coordinates[0].as_array().unwrap() {
                        // (Lat, Lon)
                        shape.push((
                            ar.as_array().unwrap()[1].as_f64().unwrap(),
                            ar.as_array().unwrap()[0].as_f64().unwrap(),
                        ));
                    }
                    let types = get_param_array(params, "type");
                    let model_autocomplete = query::autocomplete(
                        &q,
                        &pt_datasets,
                        all_data,
                        offset,
                        limit,
                        None,
                        &cnx,
                        Some(shape),
                        &types,
                    );
                    let response = model::v1::AutocompleteResponse::from(model_autocomplete);
                    render(client, response)
                })
            });
            api.get("autocomplete", |endpoint| {
                endpoint.params(|params| {
                    params.opt_typed("q", json_dsl::string());
                    dataset_param(params);
                    paginate_param(params);
                    coord_param(params, true);
                    types_param(params);
                });
                let cnx = self.es_cnx_string.clone();
                endpoint.handle(move |client, params| {
                    let q = params
                        .find("q")
                        .and_then(|val| val.as_str())
                        .unwrap_or("")
                        .to_string();
                    let pt_datasets = get_param_array(params, "pt_dataset");
                    let all_data = params
                        .find("_all_data")
                        .map_or(false, |val| val.as_bool().unwrap());
                    let offset = params
                        .find("offset")
                        .and_then(|val| val.as_u64())
                        .unwrap_or(DEFAULT_OFFSET);
                    let limit = params
                        .find("limit")
                        .and_then(|val| val.as_u64())
                        .unwrap_or(DEFAULT_LIMIT);
                    let lon = params.find("lon").and_then(|p| p.as_f64());
                    let lat = params.find("lat").and_then(|p| p.as_f64());
                    // we have already checked that if there is a lon, lat
                    // is not None so we can unwrap
                    let coord = lon.and_then(|lon| {
                        Some(Coord {
                            lon: lon,
                            lat: lat.unwrap(),
                        })
                    });

                    let types = get_param_array(params, "type");

                    let model_autocomplete = query::autocomplete(
                        &q,
                        &pt_datasets,
                        all_data,
                        offset,
                        limit,
                        coord,
                        &cnx,
                        None,
                        &types,
                    );

                    let mut response = model::v1::AutocompleteResponse::from(model_autocomplete);

                    // Optional : add distance for each feature (in meters)
                    if let Some(ref c) = &coord {
                        if let model::v1::AutocompleteResponse::Autocomplete(autocomp_resp) =
                            &mut response
                        {
                            for feature in &mut autocomp_resp.features {
                                if let ::geojson::Value::Point(p) = &feature.geometry.value {
                                    if let [mut lon, mut lat] = p.as_slice() {
                                        let feature_coord = Coord { lon: lon, lat: lat };
                                        feature.distance =
                                            Some(feature_coord.distance_to(c) as u32);
                                    }
                                }
                            }
                        }
                    }

                    render(client, response)
                })
            });
        })
    }
}
<commit_msg>create add_distance function<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io
use super::query;
use hyper::mime::Mime;
use iron::status::Status as IronStatus;
use iron::typemap::Key;
use mimir::rubber::Rubber;
use model;
use model::v1::*;
use params::{
    coord_param, dataset_param, get_param_array, paginate_param, shape_param, types_param,
};
use prometheus;
use prometheus::Encoder;
use rustless;
use rustless::server::header;
use rustless::{Api, Nesting};
use serde;
use serde_json;
use valico::json_dsl;

use navitia_model::objects::Coord;

const DEFAULT_LIMIT: u64 = 10u64;
const DEFAULT_OFFSET: u64 = 0u64;

struct Timer;
impl Key for Timer {
    type Value = prometheus::HistogramTimer;
}

lazy_static! {
    static ref HTTP_COUNTER: prometheus::CounterVec = register_counter_vec!(
        "bragi_http_requests_total",
        "Total number of HTTP requests made.",
        &["handler", "method", "status"]
    ).unwrap();
    static ref HTTP_REQ_HISTOGRAM: prometheus::HistogramVec = register_histogram_vec!(
        "bragi_http_request_duration_seconds",
        "The HTTP request latencies in seconds.",
        &["handler", "method"],
        prometheus::exponential_buckets(0.001, 1.5, 25).unwrap()
    ).unwrap();
}

fn add_distance(autocomp_resp: &mut model::Autocomplete, origin_coord: &Coord) {
    for feature in &mut autocomp_resp.features {
        if let ::geojson::Value::Point(p) = &feature.geometry.value {
            if let [mut lon, mut lat] = p.as_slice() {
                let feature_coord = Coord { lon: lon, lat: lat };
                feature.distance = Some(feature_coord.distance_to(&origin_coord) as u32);
            }
        }
    }
}

fn render<T>(
    mut client: rustless::Client,
    obj: T,
) -> Result<rustless::Client, rustless::ErrorResponse>
where
    T: serde::Serialize + model::v1::HasStatus,
{
    client.set_json_content_type();
    client.set_header(header::AccessControlAllowOrigin::Any);
    client.set_status(obj.status());
    client.text(serde_json::to_string(&obj).unwrap())
}

pub struct ApiEndPoint {
    pub es_cnx_string: String,
}

impl ApiEndPoint {
    pub fn root(&self) -> rustless::Api {
        Api::build(|api| {
            api.get("", |endpoint| {
                endpoint.handle(|client, _params| {
                    let desc = EndPoint {
                        description: "autocomplete service".to_string(),
                    };
                    render(client, desc)
                })
            });

            api.error_formatter(|error, _media| {
                let err = if error.is::<rustless::errors::Validation>() {
                    let val_err = error.downcast::<rustless::errors::Validation>().unwrap();
                    // TODO better message, we shouldn't use {:?} but access the `path`
                    // and `detail` of all errrors in val_err.reason
                    CustomError {
                        short: "validation error".to_string(),
                        long: format!("invalid arguments {:?}", val_err.reason),
                        status: IronStatus::BadRequest,
                    }
                } else {
                    CustomError {
                        short: "bad_request".to_string(),
                        long: format!("bad request, error: {}", error),
                        status: IronStatus::BadRequest,
                    }
                };
                let mut resp = rustless::Response::from(
                    err.status,
                    Box::new(serde_json::to_string(&err).unwrap()),
                );
                resp.set_json_content_type();
                Some(resp)
            });

            api.before(|client, _params| {
                let method = client.endpoint.method.to_string();

                HTTP_REQ_HISTOGRAM
                    .get_metric_with(&labels!{
                        "handler" => client.endpoint.path.path.as_str(),
                        "method" => method.as_str(),
                    })
                    .map(|timer| {
                        client.ext.insert::<Timer>(timer.start_timer());
                    })
                    .unwrap_or_else(|err| {
                        error!("impossible to get HTTP_REQ_HISTOGRAM metrics";
                               "err" => err.to_string());
                    });
                Ok(())
            });

            api.after(|client, _params| {
                let method = client.endpoint.method.to_string();
                let code = client.status().to_string();

                HTTP_COUNTER
                    .get_metric_with(&labels!{
                        "handler" => client.endpoint.path.path.as_str(),
                        "method" => method.as_str(),
                        "status" => code.as_str(),
                    })
                    .map(|counter| counter.inc())
                    .unwrap_or_else(|err| {
                        error!("impossible to get HTTP_COUNTER metrics"; "err" => err.to_string());
                    });
                client
                    .ext
                    .remove::<Timer>()
                    .map(|timer| timer.observe_duration())
                    .unwrap_or_else(|| error!("impossible to get timers from typemap"));
                Ok(())
            });
            api.mount(self.v1());
        })
    }

    fn v1(&self) -> rustless::Api {
        Api::build(|api| {
            api.mount(self.status());
            api.mount(self.autocomplete());
            api.mount(self.features());
            api.mount(self.reverse());
            api.mount(self.metrics());
        })
    }

    fn status(&self) -> rustless::Api {
        Api::build(|api| {
            api.get("status", |endpoint| {
                let cnx = self.es_cnx_string.clone();
                endpoint.handle(move |client, _params| {
                    let status = Status {
                        version: env!("CARGO_PKG_VERSION").to_string(),
                        es: cnx.to_string(),
                        status: "good".to_string(),
                    };
                    render(client, status)
                })
            });
        })
    }

    fn metrics(&self) -> rustless::Api {
        Api::build(|api| {
            api.get("metrics", |endpoint| {
                endpoint.handle(move |mut client, _params| {
                    let encoder = prometheus::TextEncoder::new();
                    let metric_familys = prometheus::gather();
                    let mut buffer = vec![];
                    encoder.encode(&metric_familys, &mut buffer).unwrap();
                    client.set_content_type(encoder.format_type().parse::<Mime>().unwrap());
                    client.text(String::from_utf8(buffer).unwrap())
                })
            });
        })
    }

    fn reverse(&self) -> rustless::Api {
        Api::build(|api| {
            api.get("reverse", |endpoint| {
                endpoint.params(|params| {
                    coord_param(params, false);
                });
                let cnx = self.es_cnx_string.clone();
                endpoint.handle(move |client, params| {
                    let coord = ::mimir::Coord::new(
                        params.find("lon").and_then(|p| p.as_f64()).unwrap(),
                        params.find("lat").and_then(|p| p.as_f64()).unwrap(),
                    );
                    let mut rubber = Rubber::new(&cnx);
                    let model_autocomplete =
                        rubber.get_address(&coord).map_err(model::BragiError::from);

                    let response = model::v1::AutocompleteResponse::from(model_autocomplete);
                    render(client, response)
                })
            });
        })
    }

    fn features(&self) -> rustless::Api {
        Api::build(|api| {
            api.get("features/:id", |endpoint| {
                endpoint.params(|params| {
                    params.opt_typed("id", json_dsl::string());
                    dataset_param(params);
                });

                let cnx = self.es_cnx_string.clone();
                endpoint.handle(move |client, params| {
                    let id = params.find("id").unwrap().as_str().unwrap();
                    let pt_datasets = get_param_array(params, "pt_dataset");
                    let all_data = params
                        .find("_all_data")
                        .map_or(false, |val| val.as_bool().unwrap());
                    let features = query::features(&pt_datasets, all_data, &cnx, &id);
                    let response = model::v1::AutocompleteResponse::from(features);
                    render(client, response)
                })
            });
        })
    }

    fn autocomplete(&self) -> rustless::Api {
        Api::build(|api| {
            api.post("autocomplete", |endpoint| {
                endpoint.params(|params| {
                    params.opt_typed("q", json_dsl::string());
                    dataset_param(params);
                    paginate_param(params);
                    shape_param(params);
                    types_param(params);
                });

                let cnx = self.es_cnx_string.clone();
                endpoint.handle(move |client, params| {
                    let q = params
                        .find("q")
                        .and_then(|val| val.as_str())
                        .unwrap_or("")
                        .to_string();
                    let pt_datasets = get_param_array(params, "pt_dataset");
                    let all_data = params
                        .find("_all_data")
                        .map_or(false, |val| val.as_bool().unwrap());
                    let offset = params
                        .find("offset")
                        .and_then(|val| val.as_u64())
                        .unwrap_or(DEFAULT_OFFSET);
                    let limit = params
                        .find("limit")
                        .and_then(|val| val.as_u64())
                        .unwrap_or(DEFAULT_LIMIT);
                    let geometry = params.find_path(&["shape", "geometry"]).unwrap();
                    let coordinates = geometry
                        .find_path(&["coordinates"])
                        .unwrap()
                        .as_array()
                        .unwrap();
                    let mut shape = Vec::new();
                    for ar in coordinates[0].as_array().unwrap() {
                        // (Lat, Lon)
                        shape.push((
                            ar.as_array().unwrap()[1].as_f64().unwrap(),
                            ar.as_array().unwrap()[0].as_f64().unwrap(),
                        ));
                    }
                    let types = get_param_array(params, "type");
                    let model_autocomplete = query::autocomplete(
                        &q,
                        &pt_datasets,
                        all_data,
                        offset,
                        limit,
                        None,
                        &cnx,
                        Some(shape),
                        &types,
                    );
                    let response = model::v1::AutocompleteResponse::from(model_autocomplete);
                    render(client, response)
                })
            });
            api.get("autocomplete", |endpoint| {
                endpoint.params(|params| {
                    params.opt_typed("q", json_dsl::string());
                    dataset_param(params);
                    paginate_param(params);
                    coord_param(params, true);
                    types_param(params);
                });
                let cnx = self.es_cnx_string.clone();
                endpoint.handle(move |client, params| {
                    let q = params
                        .find("q")
                        .and_then(|val| val.as_str())
                        .unwrap_or("")
                        .to_string();
                    let pt_datasets = get_param_array(params, "pt_dataset");
                    let all_data = params
                        .find("_all_data")
                        .map_or(false, |val| val.as_bool().unwrap());
                    let offset = params
                        .find("offset")
                        .and_then(|val| val.as_u64())
                        .unwrap_or(DEFAULT_OFFSET);
                    let limit = params
                        .find("limit")
                        .and_then(|val| val.as_u64())
                        .unwrap_or(DEFAULT_LIMIT);
                    let lon = params.find("lon").and_then(|p| p.as_f64());
                    let lat = params.find("lat").and_then(|p| p.as_f64());
                    // we have already checked that if there is a lon, lat
                    // is not None so we can unwrap
                    let coord = lon.and_then(|lon| {
                        Some(Coord {
                            lon: lon,
                            lat: lat.unwrap(),
                        })
                    });

                    let types = get_param_array(params, "type");

                    let model_autocomplete = query::autocomplete(
                        &q,
                        &pt_datasets,
                        all_data,
                        offset,
                        limit,
                        coord,
                        &cnx,
                        None,
                        &types,
                    );

                    let mut response = model::v1::AutocompleteResponse::from(model_autocomplete);

                    // Optional : add distance for each feature (in meters)
                    if let Some(origin_coord) = &coord {
                        if let model::v1::AutocompleteResponse::Autocomplete(autocomplete_resp) =
                            &mut response
                        {
                            add_distance(autocomplete_resp, origin_coord);
                        }
                    }

                    render(client, response)
                })
            });
        })
    }
}
<|endoftext|>


<commit_before>use crate::consts::{KeyNumber, FULL_KEY_LEN, KDFALG, PKGALG, PUBLIC_KEY_LEN, SIG_LEN};
use crate::errors::Error;

use alloc::string::String;
use core::convert::TryInto;
use core::ops::DerefMut;
use ed25519_dalek::{Digest, Sha512};
use rand_core::{CryptoRng, RngCore};
use zeroize::{Zeroize, Zeroizing};

/// The public half of a keypair.
///
/// You will need this if you are trying to verify a signature.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct PublicKey {
    pub(crate) keynum: KeyNumber,
    pub(crate) key: [u8; PUBLIC_KEY_LEN],
}

impl PublicKey {
    /// The public key's bytes.
    pub fn key(&self) -> [u8; PUBLIC_KEY_LEN] {
        self.key
    }

    /// The public key's identifying number.
    pub fn keynum(&self) -> KeyNumber {
        self.keynum
    }
}

/// Key derivation options available when creating a new key.
#[derive(Clone)]
pub enum NewKeyOpts {
    /// Don't encrypt the secret key.
    NoEncryption,
    /// Encrypt the secret key with a passphrase.
    Encrypted {
        /// Passphrase to encrypt the key with.
        passphrase: String,
        /// The number of key derivation rounds to apply to the password.
        ///
        /// If you're unsure of what this should be set to, use [the default].
        ///
        /// [the default]: crate::consts::DEFAULT_KDF_ROUNDS
        kdf_rounds: u32,
    },
}

impl Drop for NewKeyOpts {
    fn drop(&mut self) {
        if let NewKeyOpts::Encrypted { passphrase, .. } = self {
            passphrase.zeroize();
        }
    }
}

impl core::fmt::Debug for NewKeyOpts {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::NoEncryption => f.write_str("NoEncryption"),
            Self::Encrypted { kdf_rounds, .. } => f
                .debug_struct("Encrypted")
                .field("passphrase", &"<concealed>")
                .field("kdf_rounds", kdf_rounds)
                .finish(),
        }
    }
}

/// The full secret keypair.
///
/// You will need this if you want to create signatures.
#[derive(Clone)]
#[cfg_attr(test, derive(Debug, PartialEq))] // Makes the encoding tests nicer.
pub struct PrivateKey {
    pub(crate) public_key_alg: [u8; 2],
    pub(crate) kdf_alg: [u8; 2],
    pub(crate) kdf_rounds: u32,
    pub(crate) salt: [u8; 16],
    pub(crate) checksum: [u8; 8],
    pub(super) keynum: KeyNumber,
    pub(super) complete_key: Zeroizing<[u8; FULL_KEY_LEN]>,
}

impl PrivateKey {
    /// Generates a new random secret (private) key with the provided options.
    ///
    /// # Errors
    ///
    /// This only returns an error if the provided password was empty.
    pub fn generate<R: CryptoRng + RngCore>(
        rng: &mut R,
        derivation_info: NewKeyOpts,
    ) -> Result<Self, Error> {
        let keynum = KeyNumber::generate(rng);

        let key_pair = ed25519_dalek::Keypair::generate(rng);

        let mut skey = Zeroizing::new(key_pair.secret.to_bytes());
        let pkey = key_pair.public.to_bytes();

        let mut salt = [0; 16];
        rng.fill_bytes(&mut salt);

        let kdf_rounds = if let NewKeyOpts::Encrypted {
            passphrase,
            kdf_rounds,
        } = &derivation_info
        {
            let kdf_rounds = *kdf_rounds;
            Self::inner_kdf_mix(skey.deref_mut(), kdf_rounds, &salt, passphrase)?;
            kdf_rounds
        } else {
            0
        };

        let mut complete_key = Zeroizing::new([0u8; FULL_KEY_LEN]);
        complete_key[32..].copy_from_slice(&pkey);
        complete_key[..32].copy_from_slice(skey.as_ref());

        let checksum = Self::calculate_checksum(&complete_key);

        Ok(Self {
            public_key_alg: PKGALG,
            kdf_alg: KDFALG,
            kdf_rounds,
            salt,
            checksum,
            keynum,
            complete_key,
        })
    }

    /// Decrypts a secret key that was stored in encrypted form with the passphrase.
    ///
    /// # Errors
    ///
    /// This returns an error if the provided password was empty or if it failed to decrypt the key.
    ///
    /// The wrong password does not cause an error, but instead yields an incorrect key.
    pub fn decrypt_with_password(&mut self, passphrase: &str) -> Result<(), Error> {
        let mut encrypted_key = self.complete_key.clone(); // Cheap :)

        match Self::inner_kdf_mix(
            &mut encrypted_key[..],
            self.kdf_rounds,
            &self.salt,
            passphrase,
        ) {
            Ok(_) => {
                let current_checksum = Self::calculate_checksum(&encrypted_key);

                // Non-constant time is fine since checksum is public.
                if current_checksum != self.checksum {
                    return Err(Error::BadPassword);
                }

                self.complete_key = encrypted_key;
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    fn calculate_checksum(complete_key: &[u8; FULL_KEY_LEN]) -> [u8; 8] {
        let digest = Sha512::digest(complete_key);
        let mut checksum = [0; 8];
        checksum.copy_from_slice(&digest.as_ref()[0..8]);
        checksum
    }

    fn inner_kdf_mix(
        secret_key: &mut [u8],
        rounds: u32,
        salt: &[u8],
        passphrase: &str,
    ) -> Result<(), Error> {
        if rounds == 0 {
            return Ok(());
        }

        let mut xorkey = [0; FULL_KEY_LEN];

        bcrypt_pbkdf::bcrypt_pbkdf(passphrase, salt, rounds, &mut xorkey)
            .map_err(|_| Error::BadPassword)?;

        for (prv, xor) in secret_key.iter_mut().zip(xorkey.iter()) {
            *prv ^= xor;
        }

        Ok(())
    }

    /// Returns the public half of this secret keypair.
    pub fn public(&self) -> PublicKey {
        // This `unwrap()` gets erased in release mode.
        PublicKey {
            key: self.complete_key[32..].try_into().unwrap(),
            keynum: self.keynum,
        }
    }

    /// Returns if this key was stored encrypted.
    pub fn is_encrypted(&self) -> bool {
        self.kdf_rounds != 0
    }
}

/// A signature
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct Signature {
    pub(crate) keynum: KeyNumber,
    pub(crate) sig: [u8; SIG_LEN],
}

impl Signature {
    /// The ID of the keypair which created this signature.
    ///
    /// This is useful to determine if you have the right key to verify a signature.
    pub fn signer_keynum(&self) -> KeyNumber {
        self.keynum
    }

    /// Returns the signature's raw bytes.
    pub fn signature(&self) -> [u8; SIG_LEN] {
        self.sig
    }

    pub(super) fn new(keynum: KeyNumber, sig: [u8; SIG_LEN]) -> Self {
        Self { keynum, sig }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use core::fmt::Debug;
    use core::hash::Hash;
    use static_assertions::assert_impl_all;

    assert_impl_all!(
        PublicKey: Clone,
        Copy,
        Debug,
        Eq,
        PartialEq,
        Hash,
        Send,
        Sync
    );

    assert_impl_all!(PrivateKey: Clone, Send, Sync);

    assert_impl_all!(NewKeyOpts: Clone, Debug, Send, Sync);

    assert_impl_all!(
        Signature: Clone,
        Copy,
        Debug,
        Eq,
        PartialEq,
        Hash,
        Send,
        Sync
    );
}
<commit_msg>Fix key encryption bugs<commit_after>use crate::consts::{KeyNumber, FULL_KEY_LEN, KDFALG, PKGALG, PUBLIC_KEY_LEN, SIG_LEN};
use crate::errors::Error;

use alloc::string::String;
use core::convert::TryInto;
use ed25519_dalek::{Digest, Sha512};
use rand_core::{CryptoRng, RngCore};
use zeroize::{Zeroize, Zeroizing};

/// The public half of a keypair.
///
/// You will need this if you are trying to verify a signature.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct PublicKey {
    pub(crate) keynum: KeyNumber,
    pub(crate) key: [u8; PUBLIC_KEY_LEN],
}

impl PublicKey {
    /// The public key's bytes.
    pub fn key(&self) -> [u8; PUBLIC_KEY_LEN] {
        self.key
    }

    /// The public key's identifying number.
    pub fn keynum(&self) -> KeyNumber {
        self.keynum
    }
}

/// Key derivation options available when creating a new key.
#[derive(Clone)]
pub enum NewKeyOpts {
    /// Don't encrypt the secret key.
    NoEncryption,
    /// Encrypt the secret key with a passphrase.
    Encrypted {
        /// Passphrase to encrypt the key with.
        passphrase: String,
        /// The number of key derivation rounds to apply to the password.
        ///
        /// If you're unsure of what this should be set to, use [the default].
        ///
        /// [the default]: crate::consts::DEFAULT_KDF_ROUNDS
        kdf_rounds: u32,
    },
}

impl Drop for NewKeyOpts {
    fn drop(&mut self) {
        if let NewKeyOpts::Encrypted { passphrase, .. } = self {
            passphrase.zeroize();
        }
    }
}

impl core::fmt::Debug for NewKeyOpts {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::NoEncryption => f.write_str("NoEncryption"),
            Self::Encrypted { kdf_rounds, .. } => f
                .debug_struct("Encrypted")
                .field("passphrase", &"<concealed>")
                .field("kdf_rounds", kdf_rounds)
                .finish(),
        }
    }
}

struct UnencryptedKey(Zeroizing<[u8; FULL_KEY_LEN]>);

/// The full secret keypair.
///
/// You will need this if you want to create signatures.
#[derive(Clone)]
#[cfg_attr(test, derive(Debug, PartialEq))] // Makes the encoding tests nicer.
pub struct PrivateKey {
    pub(crate) public_key_alg: [u8; 2],
    pub(crate) kdf_alg: [u8; 2],
    pub(crate) kdf_rounds: u32,
    pub(crate) salt: [u8; 16],
    pub(crate) checksum: [u8; 8],
    pub(super) keynum: KeyNumber,
    pub(super) complete_key: Zeroizing<[u8; FULL_KEY_LEN]>,
}

impl PrivateKey {
    /// Generates a new random secret (private) key with the provided options.
    ///
    /// # Errors
    ///
    /// This only returns an error if the provided password was empty.
    pub fn generate<R: CryptoRng + RngCore>(
        rng: &mut R,
        derivation_info: NewKeyOpts,
    ) -> Result<Self, Error> {
        let keynum = KeyNumber::generate(rng);

        let key_pair = ed25519_dalek::Keypair::generate(rng);

        let mut complete_key = {
            let skey = Zeroizing::new(key_pair.secret.to_bytes());
            let pkey = key_pair.public.to_bytes();

            let mut complete_key = Zeroizing::new([0u8; FULL_KEY_LEN]);
            complete_key[..32].copy_from_slice(skey.as_ref());
            complete_key[32..].copy_from_slice(&pkey);
            UnencryptedKey(complete_key)
        };

        let checksum = Self::calculate_checksum(&complete_key);

        let mut salt = [0; 16];
        rng.fill_bytes(&mut salt);

        let kdf_rounds = if let NewKeyOpts::Encrypted {
            passphrase,
            kdf_rounds,
        } = &derivation_info
        {
            let kdf_rounds = *kdf_rounds;
            Self::inner_kdf_mix(&mut complete_key.0[..32], kdf_rounds, &salt, passphrase)?;
            kdf_rounds
        } else {
            0
        };

        Ok(Self {
            public_key_alg: PKGALG,
            kdf_alg: KDFALG,
            kdf_rounds,
            salt,
            checksum,
            keynum,
            // The state of the key doesn't matter at this stage.
            complete_key: complete_key.0,
        })
    }

    /// Decrypts a secret key that was stored in encrypted form with the passphrase.
    ///
    /// # Errors
    ///
    /// This returns an error if the provided password was empty or if it failed to decrypt the key.
    pub fn decrypt_with_password(&mut self, passphrase: &str) -> Result<(), Error> {
        let mut encrypted_key = self.complete_key.clone(); // Cheap :)

        match Self::inner_kdf_mix(
            &mut encrypted_key[..32],
            self.kdf_rounds,
            &self.salt,
            passphrase,
        ) {
            Ok(_) => {
                // Since the decryption worked, its now "unencrypted", even if the passphrase was wrong
                // and the value is garbage.
                let decrypted_key = UnencryptedKey(encrypted_key);
                let current_checksum = Self::calculate_checksum(&decrypted_key);

                // Non-constant time is fine since checksum is public.
                if current_checksum != self.checksum {
                    return Err(Error::BadPassword);
                }

                // Confirmed the decryption worked, mutating the key structure.
                self.complete_key = decrypted_key.0;
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    fn calculate_checksum(complete_key: &UnencryptedKey) -> [u8; 8] {
        let digest = Sha512::digest(complete_key.0.as_ref());
        let mut checksum = [0; 8];
        checksum.copy_from_slice(&digest.as_ref()[0..8]);
        checksum
    }

    fn inner_kdf_mix(
        secret_key: &mut [u8],
        rounds: u32,
        salt: &[u8],
        passphrase: &str,
    ) -> Result<(), Error> {
        if rounds == 0 {
            return Ok(());
        }

        let mut xorkey = [0; FULL_KEY_LEN];

        bcrypt_pbkdf::bcrypt_pbkdf(passphrase, salt, rounds, &mut xorkey)
            .map_err(|_| Error::BadPassword)?;

        for (prv, xor) in secret_key.iter_mut().zip(xorkey.iter()) {
            *prv ^= xor;
        }

        Ok(())
    }

    /// Returns the public half of this secret keypair.
    pub fn public(&self) -> PublicKey {
        // This `unwrap()` gets erased in release mode.
        PublicKey {
            key: self.complete_key[32..].try_into().unwrap(),
            keynum: self.keynum,
        }
    }

    /// Returns if this key was stored encrypted.
    pub fn is_encrypted(&self) -> bool {
        self.kdf_rounds != 0
    }
}

/// A signature
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct Signature {
    pub(crate) keynum: KeyNumber,
    pub(crate) sig: [u8; SIG_LEN],
}

impl Signature {
    /// The ID of the keypair which created this signature.
    ///
    /// This is useful to determine if you have the right key to verify a signature.
    pub fn signer_keynum(&self) -> KeyNumber {
        self.keynum
    }

    /// Returns the signature's raw bytes.
    pub fn signature(&self) -> [u8; SIG_LEN] {
        self.sig
    }

    pub(super) fn new(keynum: KeyNumber, sig: [u8; SIG_LEN]) -> Self {
        Self { keynum, sig }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use core::fmt::Debug;
    use core::hash::Hash;
    use static_assertions::assert_impl_all;

    assert_impl_all!(
        PublicKey: Clone,
        Copy,
        Debug,
        Eq,
        PartialEq,
        Hash,
        Send,
        Sync
    );

    assert_impl_all!(PrivateKey: Clone, Send, Sync);

    assert_impl_all!(NewKeyOpts: Clone, Debug, Send, Sync);

    assert_impl_all!(
        Signature: Clone,
        Copy,
        Debug,
        Eq,
        PartialEq,
        Hash,
        Send,
        Sync
    );
}
<|endoftext|>


<commit_before>mod score;
pub mod consts;

pub use self::score::*;

/// Compares two characters case-insensitively
///
/// # Examples
///
/// ```
/// assert!(rff::fuzzy::eq('a', 'A'));
/// ```
pub fn eq(a: char, b: char) -> bool {
    a.to_uppercase().eq(b.to_uppercase())
}

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

    #[test]
    fn test_eq() {
        assert!(eq('a', 'A'));
        assert!(eq('山', '山'));
        assert!(!eq('a', 'b'));
    }

    #[bench]
    fn bench_eq_same(b: &mut Bencher) {
        b.iter(|| eq('a', 'a'));
    }

    #[bench]
    fn bench_eq_ascii(b: &mut Bencher) {
        b.iter(|| eq('a', 'A'));
    }

    #[bench]
    fn bench_eq_utf8(b: &mut Bencher) {
        b.iter(|| eq('ø', 'Ø'));
    }
}
<commit_msg>Optimize fuzzy::eq(), add mixed benchmark<commit_after>mod score;
pub mod consts;

use std::ascii::AsciiExt;
pub use self::score::*;

/// Compares two characters case-insensitively
///
/// # Examples
///
/// ```
/// assert!(rff::fuzzy::eq('a', 'A'));
/// ```
pub fn eq(a: char, b: char) -> bool {
    match a {
        a if a == b => true,
        a if a.is_ascii() && !b.is_ascii() => false,
        a if !a.is_ascii() && b.is_ascii() => false,
        a if a.is_ascii() && b.is_ascii() => {
            a.to_ascii_lowercase().eq(&b.to_ascii_lowercase())
        },
        a => a.to_lowercase().eq(b.to_lowercase())
    }
}

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

    #[test]
    fn test_eq() {
        assert!(eq('a', 'A'));
        assert!(eq('山', '山'));
        assert!(!eq('a', 'b'));
    }

    #[bench]
    fn bench_eq_same(b: &mut Bencher) {
        b.iter(|| eq('a', 'a'));
    }

    #[bench]
    fn bench_eq_ascii(b: &mut Bencher) {
        b.iter(|| eq('a', 'A'));
    }

    #[bench]
    fn bench_eq_utf8(b: &mut Bencher) {
        b.iter(|| eq('ø', 'Ø'));
    }

    #[bench]
    fn bench_eq_mixed(b: &mut Bencher) {
        b.iter(|| eq('a', 'Ø'));
    }
}
<|endoftext|>