<commit_before>// See LICENSE file for copyright and license details.

use extra::json;
use collections::hashmap::HashMap;
use serialize::Decodable;
use glfw;
use gl;
use cgmath::vector::{
    Vec3,
    Vec2,
};
use glh = gl_helpers;
use camera::Camera;
use map::MapPosIter;
use geom::Geom;
use tile_picker::TilePicker;
use obj;
use mesh::Mesh;
use misc::read_file;
use core::{
    Core,
    Unit,
    CommandEndTurn,
    CommandMove,
    CommandCreateUnit,
    CommandAttackUnit,
    Event,
    EventMove,
    EventEndTurn,
    EventCreateUnit,
    EventAttackUnit,
};
use core_types::{
    Size2,
    Int,
    Bool,
    UnitId,
    PlayerId,
    MapPos,
};
use gl_types::{
    Scene,
    VertexCoord,
    TextureCoord,
    Float,
    Point2,
    MatId,
    ShaderId,
};
use event_visualizer::{
    EventVisualizer,
    EventMoveVisualizer,
    EventEndTurnVisualizer,
    EventCreateUnitVisualizer,
    EventAttackUnitVisualizer,
};
use game_state::GameState;
use pathfinder::Pathfinder;

fn build_hex_mesh(&geom: &Geom, map_size: Size2<Int>) -> ~[VertexCoord] {
    let mut vertex_data = ~[];
    for tile_pos in MapPosIter::new(map_size) {
        let pos = geom.map_pos_to_world_pos(tile_pos);
        for num in range(0 as Int, 6) {
            let vertex = geom.index_to_hex_vertex(num);
            let next_vertex = geom.index_to_hex_vertex(num + 1);
            vertex_data.push(pos + vertex);
            vertex_data.push(pos + next_vertex);
            vertex_data.push(pos + Vec3::zero());
        }
    }
    vertex_data
}

fn build_hex_tex_coord(map_size: Size2<Int>) -> ~[TextureCoord] {
    let mut vertex_data = ~[];
    for _ in MapPosIter::new(map_size) {
        for _ in range(0, 6) {
            vertex_data.push(Vec2{x: 0.0, y: 0.0});
            vertex_data.push(Vec2{x: 1.0, y: 0.0});
            vertex_data.push(Vec2{x: 0.5, y: 0.5});
        }
    }
    vertex_data
}

fn read_win_size(config_path: &str) -> Size2<Int> {
    let path = Path::new(config_path);
    let json = json::from_str(read_file(&path)).unwrap();
    let mut decoder = json::Decoder::new(json);
    let size: Size2<Int> = Decodable::decode(&mut decoder);
    size
}

fn init_win(win_size: Size2<Int>) -> glfw::Window {
    glfw::set_error_callback(~glfw::LogErrorHandler);
    let init_status = glfw::init();
    if !init_status.is_ok() {
        fail!("Failed to initialize GLFW");
    }
    let win = glfw::Window::create(
        win_size.x as u32,
        win_size.y as u32,
        "Marauder",
        glfw::Windowed,
    ).unwrap();
    win.make_context_current();
    win.set_all_polling(true);
    win
}

pub struct Visualizer<'a> {
    program: ShaderId,
    map_mesh: Mesh,
    unit_mesh: Mesh,
    mvp_mat: MatId,
    win: glfw::Window,
    mouse_pos: Point2<Float>,
    camera: Camera,
    picker: TilePicker,
    selected_tile_pos: Option<MapPos>,
    selected_unit_id: Option<UnitId>,
    geom: Geom,
    scenes: HashMap<PlayerId, Scene>,
    core: ~Core<'a>,
    event: Option<Event>,
    event_visualizer: Option<~EventVisualizer>,
    game_state: HashMap<PlayerId, GameState>,
    pathfinders: HashMap<PlayerId, Pathfinder>,
}

impl<'a> Visualizer<'a> {
    pub fn new() -> ~Visualizer {
        let players_count = 2;
        let win_size = read_win_size("config.json");
        let win = init_win(win_size);
        let geom = Geom::new();
        let core = Core::new();
        let map_size = core.map_size();
        let mut vis = ~Visualizer {
            program: 0,
            map_mesh: Mesh::new(),
            unit_mesh: Mesh::new(),
            mvp_mat: 0,
            win: win,
            mouse_pos: Vec2::zero(),
            camera: Camera::new(),
            picker: TilePicker::new(win_size),
            selected_tile_pos: None,
            selected_unit_id: None,
            geom: geom,
            scenes: {
                let mut m = HashMap::new();
                for i in range(0 as PlayerId, players_count) {
                    m.insert(i, HashMap::new());
                }
                m
            },
            game_state: {
                let mut m = HashMap::new();
                for i in range(0 as PlayerId, players_count) {
                    m.insert(i, GameState::new());
                }
                m
            },
            core: core,
            event_visualizer: None,
            event: None,
            pathfinders: {
                let mut m = HashMap::new();
                for i in range(0 as PlayerId, players_count) {
                    m.insert(i, Pathfinder::new(map_size));
                }
                m
            },
        };
        vis.init_opengl();
        vis.picker.init(&geom, vis.core.map_size());
        vis.init_models();
        vis
    }

    fn win<'a>(&'a self) -> &'a glfw::Window {
        &self.win
    }

    fn init_models(&mut self) {
        self.program = glh::compile_program(
            read_file(&Path::new("normal.vs.glsl")),
            read_file(&Path::new("normal.fs.glsl")),
        );
        gl::UseProgram(self.program);
        self.mvp_mat = glh::get_uniform(self.program, "mvp_mat");
        let vertex_coordinates_attr = glh::get_attr(
            self.program, "in_vertex_coordinates");
        gl::EnableVertexAttribArray(vertex_coordinates_attr);
        glh::vertex_attrib_pointer(vertex_coordinates_attr, 3);
        let texture_coords_attr = glh::get_attr(
            self.program, "in_texture_coordinates");
        gl::EnableVertexAttribArray(texture_coords_attr);
        glh::vertex_attrib_pointer(texture_coords_attr, 3);
        let map_size = self.core.map_size();
        let map_vertex_data = build_hex_mesh(&self.geom, map_size);
        self.map_mesh.set_vertex_coords(map_vertex_data);
        self.map_mesh.set_texture_coords(build_hex_tex_coord(map_size));
        self.map_mesh.set_texture(glh::load_texture(~"data/floor.png"));
        let unit_obj = obj::Model::new("data/tank.obj");
        self.unit_mesh.set_vertex_coords(unit_obj.build());
        self.unit_mesh.set_texture_coords(unit_obj.build_tex_coord());
        self.unit_mesh.set_texture(glh::load_texture(~"data/tank.png"));
    }

    fn scene<'a>(&'a self) -> &'a Scene {
        self.scenes.get(&self.core.player_id())
    }

    fn unit_at_opt(&'a self, pos: MapPos) -> Option<&'a Unit> {
        let mut res = None;
        let id = self.core.player_id();
        for (_, unit) in self.game_state.get(&id).units.iter() {
            if unit.pos == pos {
                res = Some(unit);
                break;
            }
        }
        res
    }

    fn init_opengl(&mut self) {
        gl::load_with(glfw::get_proc_address);
        gl::Enable(gl::DEPTH_TEST);
    }

    fn cleanup_opengl(&self) {
        gl::DeleteProgram(self.program);
    }

    fn draw_units(&self) {
        gl::UseProgram(self.program);
        for (_, unit) in self.scene().iter() {
            let m = glh::tr(self.camera.mat(), unit.pos);
            glh::uniform_mat4f(self.mvp_mat, &m);
            self.unit_mesh.draw(self.program);
        }
    }

    fn draw_map(&self) {
        glh::uniform_mat4f(self.mvp_mat, &self.camera.mat());
        self.map_mesh.draw(self.program);
    }

    fn draw(&mut self) {
        gl::ClearColor(0.3, 0.3, 0.3, 1.0);
        gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);
        gl::UseProgram(self.program);
        self.draw_units();
        self.draw_map();
        if !self.event_visualizer.is_none() {
            let scene = self.scenes.get_mut(&self.core.player_id());
            self.event_visualizer.get_mut_ref().draw(&self.geom, scene);
        }
        self.win().swap_buffers();
    }

    pub fn is_running(&self) -> Bool {
        return !self.win().should_close()
    }

    fn handle_key_event(&mut self, key: glfw::Key) {
        match key {
            glfw::KeyEscape | glfw::KeyQ => {
                self.win().set_should_close(true);
            },
            glfw::KeyT => {
                self.core.do_command(CommandEndTurn);
            },
            glfw::KeyU => {
                let pos_opt = self.selected_tile_pos;
                if pos_opt.is_some() {
                    let pos = pos_opt.unwrap();
                    if self.unit_at_opt(pos).is_none() {
                        let cmd = CommandCreateUnit(pos);
                        self.core.do_command(cmd);
                    }
                }
            },
            glfw::KeyA => {
                let pos_opt = self.selected_tile_pos;
                if pos_opt.is_some() {
                    let pos = pos_opt.unwrap();
                    if self.unit_at_opt(pos).is_some() {
                        let defender_id = self.unit_at_opt(pos).unwrap().id;
                        let cmd = CommandAttackUnit(0, defender_id);
                        self.core.do_command(cmd);
                    }
                }
            },
            glfw::KeySpace => println!("space"),
            glfw::KeyUp => self.camera.move(270.0),
            glfw::KeyDown => self.camera.move(90.0),
            glfw::KeyRight => self.camera.move(0.0),
            glfw::KeyLeft => self.camera.move(180.0),
            glfw::KeyMinus => self.camera.zoom += 1.0,
            glfw::KeyEqual => self.camera.zoom -= 1.0,
            _ => {},
        }
    }

    fn handle_cursor_pos_event(&mut self, pos: Point2<Float>) {
        let button = self.win().get_mouse_button(glfw::MouseButtonRight);
        if button == glfw::Press {
            self.camera.z_angle += (self.mouse_pos.x - pos.x) / 2.0;
            self.camera.x_angle += (self.mouse_pos.y - pos.y) / 2.0;
        }
        self.mouse_pos = pos;
    }

    fn handle_mouse_button_event(&mut self) {
        if self.selected_tile_pos.is_some() {
            let pos = self.selected_tile_pos.unwrap();
            if self.unit_at_opt(pos).is_some() {
                let unit_id = self.unit_at_opt(pos).unwrap().id;
                let state = self.game_state.get_mut(&self.core.player_id());
                self.selected_unit_id = Some(unit_id);
                let pf = self.pathfinders.get_mut(&self.core.player_id());
                pf.fill_map(state, state.units.get(&unit_id));
            } else if self.selected_unit_id.is_some() {
                let unit_id = self.selected_unit_id.unwrap();
                let pf = self.pathfinders.get_mut(&self.core.player_id());
                let path = pf.get_path(pos);
                self.core.do_command(CommandMove(unit_id, path));
            }
        }
    }

    fn get_events(&mut self) -> ~[glfw::WindowEvent] {
        glfw::poll_events();
        let mut events = ~[];
        for (_, event) in self.win().flush_events() {
            events.push(event);
        }
        events
    }

    fn handle_event(&mut self, event: glfw::WindowEvent) {
        match event {
            glfw::KeyEvent(key, _, glfw::Press, _) => {
                self.handle_key_event(key);
            },
            glfw::CursorPosEvent(x, y) => {
                let p = Vec2{x: x as Float, y: y as Float};
                self.handle_cursor_pos_event(p);
            },
            glfw::MouseButtonEvent(glfw::MouseButtonLeft, glfw::Press, _) => {
                self.handle_mouse_button_event();
            },
            glfw::SizeEvent(w, h) => {
                gl::Viewport(0, 0, w, h);
                self.picker.set_win_size(Size2{x: w, y: h});
            },
            _ => {},
        }
    }

    fn handle_events(&mut self) {
        for event in self.get_events().iter() {
            self.handle_event(*event);
        }
    }

    fn pick_tile(&mut self) {
        let mouse_pos = Vec2 {
            x: self.mouse_pos.x as Int,
            y: self.mouse_pos.y as Int,
        };
        self.selected_tile_pos = self.picker.pick_tile(&self.camera, mouse_pos);
    }

    pub fn make_event_visualizer(&mut self, event: &Event) -> ~EventVisualizer {
        match *event {
            EventMove(ref unit_id, ref path) => {
                EventMoveVisualizer::new(*unit_id, path.clone())
            },
            EventEndTurn(_, _) => {
                EventEndTurnVisualizer::new()
            },
            EventCreateUnit(id, ref pos) => {
                EventCreateUnitVisualizer::new(id, *pos)
            },
            EventAttackUnit(attacker_id, defender_id) => {
                EventAttackUnitVisualizer::new(attacker_id, defender_id)
            },
        }
    }

    pub fn logic(&mut self) {
        if self.event_visualizer.is_none() {
            let event_opt = self.core.get_event();
            if event_opt.is_some() {
                let event = event_opt.unwrap();
                let vis = self.make_event_visualizer(&event);
                self.event = Some(event);
                self.event_visualizer = Some(vis);
            }
        } else if self.event_visualizer.get_ref().is_finished() {
            let scene = self.scenes.get_mut(&self.core.player_id());
            self.event_visualizer.get_mut_ref().end(&self.geom, scene);
            let state = self.game_state.get_mut(&self.core.player_id());
            state.apply_event(self.event.get_ref());
            self.event_visualizer = None;
            self.event = None;
            if self.selected_unit_id.is_some() {
                let unit_id = self.selected_unit_id.unwrap();
                let pf = self.pathfinders.get_mut(&self.core.player_id());
                pf.fill_map(state, state.units.get(&unit_id));
            }
        }
    }

    pub fn tick(&mut self) {
        self.handle_events();
        self.logic();
        self.pick_tile();
        self.draw();
    }
}

impl<'a> Drop for Visualizer<'a> {
    fn drop(&mut self) {
        self.cleanup_opengl();
    }
}

// vim: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
<commit_msg>Visualizer: fn get_{scenes,game_states,pathfinders}<commit_after>// See LICENSE file for copyright and license details.

use extra::json;
use collections::hashmap::HashMap;
use serialize::Decodable;
use glfw;
use gl;
use cgmath::vector::{
    Vec3,
    Vec2,
};
use glh = gl_helpers;
use camera::Camera;
use map::MapPosIter;
use geom::Geom;
use tile_picker::TilePicker;
use obj;
use mesh::Mesh;
use misc::read_file;
use core::{
    Core,
    Unit,
    CommandEndTurn,
    CommandMove,
    CommandCreateUnit,
    CommandAttackUnit,
    Event,
    EventMove,
    EventEndTurn,
    EventCreateUnit,
    EventAttackUnit,
};
use core_types::{
    Size2,
    Int,
    Bool,
    UnitId,
    PlayerId,
    MapPos,
};
use gl_types::{
    Scene,
    VertexCoord,
    TextureCoord,
    Float,
    Point2,
    MatId,
    ShaderId,
};
use event_visualizer::{
    EventVisualizer,
    EventMoveVisualizer,
    EventEndTurnVisualizer,
    EventCreateUnitVisualizer,
    EventAttackUnitVisualizer,
};
use game_state::GameState;
use pathfinder::Pathfinder;

fn build_hex_mesh(&geom: &Geom, map_size: Size2<Int>) -> ~[VertexCoord] {
    let mut vertex_data = ~[];
    for tile_pos in MapPosIter::new(map_size) {
        let pos = geom.map_pos_to_world_pos(tile_pos);
        for num in range(0 as Int, 6) {
            let vertex = geom.index_to_hex_vertex(num);
            let next_vertex = geom.index_to_hex_vertex(num + 1);
            vertex_data.push(pos + vertex);
            vertex_data.push(pos + next_vertex);
            vertex_data.push(pos + Vec3::zero());
        }
    }
    vertex_data
}

fn build_hex_tex_coord(map_size: Size2<Int>) -> ~[TextureCoord] {
    let mut vertex_data = ~[];
    for _ in MapPosIter::new(map_size) {
        for _ in range(0, 6) {
            vertex_data.push(Vec2{x: 0.0, y: 0.0});
            vertex_data.push(Vec2{x: 1.0, y: 0.0});
            vertex_data.push(Vec2{x: 0.5, y: 0.5});
        }
    }
    vertex_data
}

fn read_win_size(config_path: &str) -> Size2<Int> {
    let path = Path::new(config_path);
    let json = json::from_str(read_file(&path)).unwrap();
    let mut decoder = json::Decoder::new(json);
    let size: Size2<Int> = Decodable::decode(&mut decoder);
    size
}

fn init_win(win_size: Size2<Int>) -> glfw::Window {
    glfw::set_error_callback(~glfw::LogErrorHandler);
    let init_status = glfw::init();
    if !init_status.is_ok() {
        fail!("Failed to initialize GLFW");
    }
    let win = glfw::Window::create(
        win_size.x as u32,
        win_size.y as u32,
        "Marauder",
        glfw::Windowed,
    ).unwrap();
    win.make_context_current();
    win.set_all_polling(true);
    win
}

fn get_scenes(players_count: Int) -> HashMap<PlayerId, Scene> {
    let mut m = HashMap::new();
    for i in range(0 as PlayerId, players_count) {
        m.insert(i, HashMap::new());
    }
    m
}

fn get_game_states(players_count: Int) -> HashMap<PlayerId, GameState> {
    let mut m = HashMap::new();
    for i in range(0 as PlayerId, players_count) {
        m.insert(i, GameState::new());
    }
    m
}

fn get_pathfinders(
    players_count: Int,
    map_size: Size2<Int>,
) -> HashMap<PlayerId, Pathfinder> {
    let mut m = HashMap::new();
    for i in range(0 as PlayerId, players_count) {
        m.insert(i, Pathfinder::new(map_size));
    }
    m
}

pub struct Visualizer<'a> {
    program: ShaderId,
    map_mesh: Mesh,
    unit_mesh: Mesh,
    mvp_mat: MatId,
    win: glfw::Window,
    mouse_pos: Point2<Float>,
    camera: Camera,
    picker: TilePicker,
    selected_tile_pos: Option<MapPos>,
    selected_unit_id: Option<UnitId>,
    geom: Geom,
    scenes: HashMap<PlayerId, Scene>,
    core: ~Core<'a>,
    event: Option<Event>,
    event_visualizer: Option<~EventVisualizer>,
    game_state: HashMap<PlayerId, GameState>,
    pathfinders: HashMap<PlayerId, Pathfinder>,
}

impl<'a> Visualizer<'a> {
    pub fn new() -> ~Visualizer {
        let players_count = 2;
        let win_size = read_win_size("config.json");
        let win = init_win(win_size);
        let geom = Geom::new();
        let core = Core::new();
        let map_size = core.map_size();
        let mut vis = ~Visualizer {
            program: 0,
            map_mesh: Mesh::new(),
            unit_mesh: Mesh::new(),
            mvp_mat: 0,
            win: win,
            mouse_pos: Vec2::zero(),
            camera: Camera::new(),
            picker: TilePicker::new(win_size),
            selected_tile_pos: None,
            selected_unit_id: None,
            geom: geom,
            core: core,
            event_visualizer: None,
            event: None,
            scenes: get_scenes(players_count),
            game_state: get_game_states(players_count),
            pathfinders: get_pathfinders(players_count, map_size),
        };
        vis.init_opengl();
        vis.picker.init(&geom, vis.core.map_size());
        vis.init_models();
        vis
    }

    fn win<'a>(&'a self) -> &'a glfw::Window {
        &self.win
    }

    fn init_models(&mut self) {
        self.program = glh::compile_program(
            read_file(&Path::new("normal.vs.glsl")),
            read_file(&Path::new("normal.fs.glsl")),
        );
        gl::UseProgram(self.program);
        self.mvp_mat = glh::get_uniform(self.program, "mvp_mat");
        let vertex_coordinates_attr = glh::get_attr(
            self.program, "in_vertex_coordinates");
        gl::EnableVertexAttribArray(vertex_coordinates_attr);
        glh::vertex_attrib_pointer(vertex_coordinates_attr, 3);
        let texture_coords_attr = glh::get_attr(
            self.program, "in_texture_coordinates");
        gl::EnableVertexAttribArray(texture_coords_attr);
        glh::vertex_attrib_pointer(texture_coords_attr, 3);
        let map_size = self.core.map_size();
        let map_vertex_data = build_hex_mesh(&self.geom, map_size);
        self.map_mesh.set_vertex_coords(map_vertex_data);
        self.map_mesh.set_texture_coords(build_hex_tex_coord(map_size));
        self.map_mesh.set_texture(glh::load_texture(~"data/floor.png"));
        let unit_obj = obj::Model::new("data/tank.obj");
        self.unit_mesh.set_vertex_coords(unit_obj.build());
        self.unit_mesh.set_texture_coords(unit_obj.build_tex_coord());
        self.unit_mesh.set_texture(glh::load_texture(~"data/tank.png"));
    }

    fn scene<'a>(&'a self) -> &'a Scene {
        self.scenes.get(&self.core.player_id())
    }

    fn unit_at_opt(&'a self, pos: MapPos) -> Option<&'a Unit> {
        let mut res = None;
        let id = self.core.player_id();
        for (_, unit) in self.game_state.get(&id).units.iter() {
            if unit.pos == pos {
                res = Some(unit);
                break;
            }
        }
        res
    }

    fn init_opengl(&mut self) {
        gl::load_with(glfw::get_proc_address);
        gl::Enable(gl::DEPTH_TEST);
    }

    fn cleanup_opengl(&self) {
        gl::DeleteProgram(self.program);
    }

    fn draw_units(&self) {
        gl::UseProgram(self.program);
        for (_, unit) in self.scene().iter() {
            let m = glh::tr(self.camera.mat(), unit.pos);
            glh::uniform_mat4f(self.mvp_mat, &m);
            self.unit_mesh.draw(self.program);
        }
    }

    fn draw_map(&self) {
        glh::uniform_mat4f(self.mvp_mat, &self.camera.mat());
        self.map_mesh.draw(self.program);
    }

    fn draw(&mut self) {
        gl::ClearColor(0.3, 0.3, 0.3, 1.0);
        gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);
        gl::UseProgram(self.program);
        self.draw_units();
        self.draw_map();
        if !self.event_visualizer.is_none() {
            let scene = self.scenes.get_mut(&self.core.player_id());
            self.event_visualizer.get_mut_ref().draw(&self.geom, scene);
        }
        self.win().swap_buffers();
    }

    pub fn is_running(&self) -> Bool {
        return !self.win().should_close()
    }

    fn handle_key_event(&mut self, key: glfw::Key) {
        match key {
            glfw::KeyEscape | glfw::KeyQ => {
                self.win().set_should_close(true);
            },
            glfw::KeyT => {
                self.core.do_command(CommandEndTurn);
            },
            glfw::KeyU => {
                let pos_opt = self.selected_tile_pos;
                if pos_opt.is_some() {
                    let pos = pos_opt.unwrap();
                    if self.unit_at_opt(pos).is_none() {
                        let cmd = CommandCreateUnit(pos);
                        self.core.do_command(cmd);
                    }
                }
            },
            glfw::KeyA => {
                let pos_opt = self.selected_tile_pos;
                if pos_opt.is_some() {
                    let pos = pos_opt.unwrap();
                    if self.unit_at_opt(pos).is_some() {
                        let defender_id = self.unit_at_opt(pos).unwrap().id;
                        let cmd = CommandAttackUnit(0, defender_id);
                        self.core.do_command(cmd);
                    }
                }
            },
            glfw::KeySpace => println!("space"),
            glfw::KeyUp => self.camera.move(270.0),
            glfw::KeyDown => self.camera.move(90.0),
            glfw::KeyRight => self.camera.move(0.0),
            glfw::KeyLeft => self.camera.move(180.0),
            glfw::KeyMinus => self.camera.zoom += 1.0,
            glfw::KeyEqual => self.camera.zoom -= 1.0,
            _ => {},
        }
    }

    fn handle_cursor_pos_event(&mut self, pos: Point2<Float>) {
        let button = self.win().get_mouse_button(glfw::MouseButtonRight);
        if button == glfw::Press {
            self.camera.z_angle += (self.mouse_pos.x - pos.x) / 2.0;
            self.camera.x_angle += (self.mouse_pos.y - pos.y) / 2.0;
        }
        self.mouse_pos = pos;
    }

    fn handle_mouse_button_event(&mut self) {
        if self.selected_tile_pos.is_some() {
            let pos = self.selected_tile_pos.unwrap();
            if self.unit_at_opt(pos).is_some() {
                let unit_id = self.unit_at_opt(pos).unwrap().id;
                let state = self.game_state.get_mut(&self.core.player_id());
                self.selected_unit_id = Some(unit_id);
                let pf = self.pathfinders.get_mut(&self.core.player_id());
                pf.fill_map(state, state.units.get(&unit_id));
            } else if self.selected_unit_id.is_some() {
                let unit_id = self.selected_unit_id.unwrap();
                let pf = self.pathfinders.get_mut(&self.core.player_id());
                let path = pf.get_path(pos);
                self.core.do_command(CommandMove(unit_id, path));
            }
        }
    }

    fn get_events(&mut self) -> ~[glfw::WindowEvent] {
        glfw::poll_events();
        let mut events = ~[];
        for (_, event) in self.win().flush_events() {
            events.push(event);
        }
        events
    }

    fn handle_event(&mut self, event: glfw::WindowEvent) {
        match event {
            glfw::KeyEvent(key, _, glfw::Press, _) => {
                self.handle_key_event(key);
            },
            glfw::CursorPosEvent(x, y) => {
                let p = Vec2{x: x as Float, y: y as Float};
                self.handle_cursor_pos_event(p);
            },
            glfw::MouseButtonEvent(glfw::MouseButtonLeft, glfw::Press, _) => {
                self.handle_mouse_button_event();
            },
            glfw::SizeEvent(w, h) => {
                gl::Viewport(0, 0, w, h);
                self.picker.set_win_size(Size2{x: w, y: h});
            },
            _ => {},
        }
    }

    fn handle_events(&mut self) {
        for event in self.get_events().iter() {
            self.handle_event(*event);
        }
    }

    fn pick_tile(&mut self) {
        let mouse_pos = Vec2 {
            x: self.mouse_pos.x as Int,
            y: self.mouse_pos.y as Int,
        };
        self.selected_tile_pos = self.picker.pick_tile(&self.camera, mouse_pos);
    }

    pub fn make_event_visualizer(&mut self, event: &Event) -> ~EventVisualizer {
        match *event {
            EventMove(ref unit_id, ref path) => {
                EventMoveVisualizer::new(*unit_id, path.clone())
            },
            EventEndTurn(_, _) => {
                EventEndTurnVisualizer::new()
            },
            EventCreateUnit(id, ref pos) => {
                EventCreateUnitVisualizer::new(id, *pos)
            },
            EventAttackUnit(attacker_id, defender_id) => {
                EventAttackUnitVisualizer::new(attacker_id, defender_id)
            },
        }
    }

    pub fn logic(&mut self) {
        if self.event_visualizer.is_none() {
            let event_opt = self.core.get_event();
            if event_opt.is_some() {
                let event = event_opt.unwrap();
                let vis = self.make_event_visualizer(&event);
                self.event = Some(event);
                self.event_visualizer = Some(vis);
            }
        } else if self.event_visualizer.get_ref().is_finished() {
            let scene = self.scenes.get_mut(&self.core.player_id());
            self.event_visualizer.get_mut_ref().end(&self.geom, scene);
            let state = self.game_state.get_mut(&self.core.player_id());
            state.apply_event(self.event.get_ref());
            self.event_visualizer = None;
            self.event = None;
            if self.selected_unit_id.is_some() {
                let unit_id = self.selected_unit_id.unwrap();
                let pf = self.pathfinders.get_mut(&self.core.player_id());
                pf.fill_map(state, state.units.get(&unit_id));
            }
        }
    }

    pub fn tick(&mut self) {
        self.handle_events();
        self.logic();
        self.pick_tile();
        self.draw();
    }
}

impl<'a> Drop for Visualizer<'a> {
    fn drop(&mut self) {
        self.cleanup_opengl();
    }
}

// vim: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
<|endoftext|>


<commit_before>use crate::errors::ParseError;
use crate::parsing_utils::{decode_str, decode_utf16, decode_windows1252, le_i32};

#[derive(Debug, Clone, PartialEq)]
pub struct CoreParser<'a> {
    data: &'a [u8],

    /// Current offset in regards to the whole view of the replay
    col: i32,
}

impl<'a> CoreParser<'a> {
    pub fn new(data: &'a [u8]) -> Self {
        CoreParser { data, col: 0 }
    }

    pub fn bytes_read(&self) -> i32 {
        self.col
    }

    /// Used for skipping some amount of data
    pub fn advance(&mut self, ind: usize) {
        self.col += ind as i32;
        self.data = &self.data[ind..];
    }

    /// Returns a slice of the replay after ensuring there is enough space for the requested slice
    pub fn view_data(&self, size: usize) -> Result<&'a [u8], ParseError> {
        if size > self.data.len() {
            Err(ParseError::InsufficientData(
                size as i32,
                self.data.len() as i32,
            ))
        } else {
            Ok(&self.data[..size])
        }
    }

    /// Take the next `size` of bytes and interpret them in an infallible fashion
    #[inline]
    pub fn take<F, T>(&mut self, size: usize, mut f: F) -> Result<T, ParseError>
    where
        F: FnMut(&'a [u8]) -> T,
    {
        let res = f(self.view_data(size)?);
        self.advance(size);
        Ok(res)
    }

    /// Take the next `size` of bytes and interpret them, but this interpretation can fail
    pub fn take_res<F, T>(&mut self, size: usize, mut f: F) -> Result<T, ParseError>
    where
        F: FnMut(&'a [u8]) -> Result<T, ParseError>,
    {
        let res = f(self.view_data(size)?)?;
        self.advance(size);
        Ok(res)
    }

    /// Repeatedly parse the same elements from replay until `size` elements parsed
    pub fn repeat<F, T>(size: usize, mut f: F) -> Result<Vec<T>, ParseError>
    where
        F: FnMut() -> Result<T, ParseError>,
    {
        if size > 25_000 {
            return Err(ParseError::ListTooLarge(size));
        }

        let mut res = Vec::with_capacity(size);
        for _ in 0..size {
            res.push(f()?);
        }
        Ok(res)
    }

    pub fn list_of<F, T>(&mut self, mut f: F) -> Result<Vec<T>, ParseError>
    where
        F: FnMut(&mut Self) -> Result<T, ParseError>,
    {
        let size = self.take(4, le_i32)?;
        CoreParser::repeat(size as usize, || f(self))
    }

    pub fn text_list(&mut self) -> Result<Vec<String>, ParseError> {
        self.list_of(CoreParser::parse_text)
    }

    /// Parses UTF-8 string from replay
    pub fn parse_str(&mut self) -> Result<&'a str, ParseError> {
        let mut size = self.take(4, le_i32)? as usize;

        // Replay 6688 has a property name that is listed as having a length of 0x5000000, but it's
        // really the `\0\0\0None` property. I'm guess at some point in Rocket League, this was a
        // bug that was fixed. What's interesting is that I couldn't find this constant in
        // `RocketLeagueReplayParser`, only rattletrap.
        if size == 0x5_000_000 {
            size = 8;
        }
        self.take_res(size, decode_str)
    }

    /// Parses either UTF-16 or Windows-1252 encoded strings
    pub fn parse_text(&mut self) -> Result<String, ParseError> {
        // The number of bytes that the string is composed of. If negative, the string is UTF-16,
        // else the string is windows 1252 encoded.
        let characters = self.take(4, le_i32)?;

        // size.abs() will panic at min_value, so we eschew it for manual checking
        if characters == 0 {
            Err(ParseError::ZeroSize)
        } else if characters > 10_000 || characters < -10_000 {
            Err(ParseError::TextTooLarge(characters))
        } else if characters < 0 {
            // We're dealing with UTF-16 and each character is two bytes, we
            // multiply the size by 2. The last two bytes included in the count are
            // null terminators
            let size = characters * -2;
            self.take_res(size as usize, |d| decode_utf16(d))
        } else {
            self.take_res(characters as usize, |d| decode_windows1252(d))
        }
    }
}

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

    #[test]
    fn parse_text_encoding() {
        // dd skip=16 count=28 if=rumble.replay of=text.replay bs=1
        let data = include_bytes!("../assets/replays/partial/text.replay");
        let mut parser = CoreParser::new(&data[..]);
        assert_eq!(parser.parse_str().unwrap(), "TAGame.Replay_Soccar_TA");
    }

    #[test]
    fn parse_text_encoding_bad() {
        // dd skip=16 count=28 if=rumble.replay of=text.replay bs=1
        let data = include_bytes!("../assets/replays/partial/text.replay");
        let mut parser = CoreParser::new(&data[..data.len() - 1]);
        let res = parser.parse_str();
        assert!(res.is_err());
        let error = res.unwrap_err();
        assert_eq!(error, ParseError::InsufficientData(24, 23));
    }

    #[test]
    fn parse_text_zero_size() {
        let mut parser = CoreParser::new(&[0, 0, 0, 0, 0]);
        let res = parser.parse_str();
        assert!(res.is_err());
        let error = res.unwrap_err();
        assert_eq!(error, ParseError::ZeroSize);
    }

    #[test]
    fn parse_text_encoding_bad_2() {
        // Test for when there is not enough data to decode text length
        // dd skip=16 count=28 if=rumble.replay of=text.replay bs=1
        let data = include_bytes!("../assets/replays/partial/text.replay");
        let mut parser = CoreParser::new(&data[..2]);
        let res = parser.parse_str();
        assert!(res.is_err());
        let error = res.unwrap_err();
        assert_eq!(error, ParseError::InsufficientData(4, 2));
    }

    #[test]
    fn parse_utf16_string() {
        // dd skip=((0x120)) count=28 if=utf-16.replay of=utf-16-text.replay bs=1
        let data = include_bytes!("../assets/replays/partial/utf-16-text.replay");
        let mut parser = CoreParser::new(&data[..]);
        let res = parser.parse_text().unwrap();
        assert_eq!(res, "\u{2623}D[e]!v1zz\u{2623}");
    }

    #[test]
    fn test_windows1252_string() {
        let data = include_bytes!("../assets/replays/partial/windows_1252.replay");
        let mut parser = CoreParser::new(&data[0x1ad..0x1c4]);
        let res = parser.parse_text().unwrap();
        assert_eq!(res, "caudillman6000\u{b3}(2)");
    }

    /// Define behavior on invalid UTF-16 sequences.
    #[test]
    fn parse_invalid_utf16_string() {
        let data = [0xfd, 0xff, 0xff, 0xff, 0xd8, 0xd8, 0x00, 0x00, 0x00, 0x00];
        let mut parser = CoreParser::new(&data[..]);
        let res = parser.parse_text().unwrap();
        assert_eq!(res, "�\u{0}");
    }
}
<commit_msg>Add couple unhappy path test cases for parsing test<commit_after>use crate::errors::ParseError;
use crate::parsing_utils::{decode_str, decode_utf16, decode_windows1252, le_i32};

#[derive(Debug, Clone, PartialEq)]
pub struct CoreParser<'a> {
    data: &'a [u8],

    /// Current offset in regards to the whole view of the replay
    col: i32,
}

impl<'a> CoreParser<'a> {
    pub fn new(data: &'a [u8]) -> Self {
        CoreParser { data, col: 0 }
    }

    pub fn bytes_read(&self) -> i32 {
        self.col
    }

    /// Used for skipping some amount of data
    pub fn advance(&mut self, ind: usize) {
        self.col += ind as i32;
        self.data = &self.data[ind..];
    }

    /// Returns a slice of the replay after ensuring there is enough space for the requested slice
    pub fn view_data(&self, size: usize) -> Result<&'a [u8], ParseError> {
        if size > self.data.len() {
            Err(ParseError::InsufficientData(
                size as i32,
                self.data.len() as i32,
            ))
        } else {
            Ok(&self.data[..size])
        }
    }

    /// Take the next `size` of bytes and interpret them in an infallible fashion
    #[inline]
    pub fn take<F, T>(&mut self, size: usize, mut f: F) -> Result<T, ParseError>
    where
        F: FnMut(&'a [u8]) -> T,
    {
        let res = f(self.view_data(size)?);
        self.advance(size);
        Ok(res)
    }

    /// Take the next `size` of bytes and interpret them, but this interpretation can fail
    pub fn take_res<F, T>(&mut self, size: usize, mut f: F) -> Result<T, ParseError>
    where
        F: FnMut(&'a [u8]) -> Result<T, ParseError>,
    {
        let res = f(self.view_data(size)?)?;
        self.advance(size);
        Ok(res)
    }

    /// Repeatedly parse the same elements from replay until `size` elements parsed
    pub fn repeat<F, T>(size: usize, mut f: F) -> Result<Vec<T>, ParseError>
    where
        F: FnMut() -> Result<T, ParseError>,
    {
        if size > 25_000 {
            return Err(ParseError::ListTooLarge(size));
        }

        let mut res = Vec::with_capacity(size);
        for _ in 0..size {
            res.push(f()?);
        }
        Ok(res)
    }

    pub fn list_of<F, T>(&mut self, mut f: F) -> Result<Vec<T>, ParseError>
    where
        F: FnMut(&mut Self) -> Result<T, ParseError>,
    {
        let size = self.take(4, le_i32)?;
        CoreParser::repeat(size as usize, || f(self))
    }

    pub fn text_list(&mut self) -> Result<Vec<String>, ParseError> {
        self.list_of(CoreParser::parse_text)
    }

    /// Parses UTF-8 string from replay
    pub fn parse_str(&mut self) -> Result<&'a str, ParseError> {
        let mut size = self.take(4, le_i32)? as usize;

        // Replay 6688 has a property name that is listed as having a length of 0x5000000, but it's
        // really the `\0\0\0None` property. I'm guess at some point in Rocket League, this was a
        // bug that was fixed. What's interesting is that I couldn't find this constant in
        // `RocketLeagueReplayParser`, only rattletrap.
        if size == 0x5_000_000 {
            size = 8;
        }
        self.take_res(size, decode_str)
    }

    /// Parses either UTF-16 or Windows-1252 encoded strings
    pub fn parse_text(&mut self) -> Result<String, ParseError> {
        // The number of bytes that the string is composed of. If negative, the string is UTF-16,
        // else the string is windows 1252 encoded.
        let characters = self.take(4, le_i32)?;

        // size.abs() will panic at min_value, so we eschew it for manual checking
        if characters == 0 {
            Err(ParseError::ZeroSize)
        } else if characters > 10_000 || characters < -10_000 {
            Err(ParseError::TextTooLarge(characters))
        } else if characters < 0 {
            // We're dealing with UTF-16 and each character is two bytes, we
            // multiply the size by 2. The last two bytes included in the count are
            // null terminators
            let size = characters * -2;
            self.take_res(size as usize, |d| decode_utf16(d))
        } else {
            self.take_res(characters as usize, |d| decode_windows1252(d))
        }
    }
}

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

    #[test]
    fn parse_text_encoding() {
        // dd skip=16 count=28 if=rumble.replay of=text.replay bs=1
        let data = include_bytes!("../assets/replays/partial/text.replay");
        let mut parser = CoreParser::new(&data[..]);
        assert_eq!(parser.parse_str().unwrap(), "TAGame.Replay_Soccar_TA");
    }

    #[test]
    fn parse_text_encoding_bad() {
        // dd skip=16 count=28 if=rumble.replay of=text.replay bs=1
        let data = include_bytes!("../assets/replays/partial/text.replay");
        let mut parser = CoreParser::new(&data[..data.len() - 1]);
        let res = parser.parse_str();
        assert!(res.is_err());
        let error = res.unwrap_err();
        assert_eq!(error, ParseError::InsufficientData(24, 23));
    }

    #[test]
    fn parse_text_zero_size() {
        let mut parser = CoreParser::new(&[0, 0, 0, 0, 0]);
        let res = parser.parse_str();
        assert!(res.is_err());
        let error = res.unwrap_err();
        assert_eq!(error, ParseError::ZeroSize);
    }

    #[test]
    fn parse_text_zero_size2() {
        let mut parser = CoreParser::new(&[0, 0, 0, 0, 0]);
        let res = parser.parse_text();
        assert!(res.is_err());
        let error = res.unwrap_err();
        assert_eq!(error, ParseError::ZeroSize);
    }

    #[test]
    fn parse_text_too_large() {
        let mut parser = CoreParser::new(&[0xcc, 0xcc, 0xcc, 0xcc, 0xcc]);
        let res = parser.parse_text();
        assert!(res.is_err());
        let error = res.unwrap_err();
        assert_eq!(error, ParseError::TextTooLarge(-858993460));
    }

    #[test]
    fn parse_text_encoding_bad_2() {
        // Test for when there is not enough data to decode text length
        // dd skip=16 count=28 if=rumble.replay of=text.replay bs=1
        let data = include_bytes!("../assets/replays/partial/text.replay");
        let mut parser = CoreParser::new(&data[..2]);
        let res = parser.parse_str();
        assert!(res.is_err());
        let error = res.unwrap_err();
        assert_eq!(error, ParseError::InsufficientData(4, 2));
    }

    #[test]
    fn parse_utf16_string() {
        // dd skip=((0x120)) count=28 if=utf-16.replay of=utf-16-text.replay bs=1
        let data = include_bytes!("../assets/replays/partial/utf-16-text.replay");
        let mut parser = CoreParser::new(&data[..]);
        let res = parser.parse_text().unwrap();
        assert_eq!(res, "\u{2623}D[e]!v1zz\u{2623}");
    }

    #[test]
    fn test_windows1252_string() {
        let data = include_bytes!("../assets/replays/partial/windows_1252.replay");
        let mut parser = CoreParser::new(&data[0x1ad..0x1c4]);
        let res = parser.parse_text().unwrap();
        assert_eq!(res, "caudillman6000\u{b3}(2)");
    }

    /// Define behavior on invalid UTF-16 sequences.
    #[test]
    fn parse_invalid_utf16_string() {
        let data = [0xfd, 0xff, 0xff, 0xff, 0xd8, 0xd8, 0x00, 0x00, 0x00, 0x00];
        let mut parser = CoreParser::new(&data[..]);
        let res = parser.parse_text().unwrap();
        assert_eq!(res, "�\u{0}");
    }
}
<|endoftext|>


<commit_before>//! Aead Ciphers

use crate::crypto::cipher::{CipherCategory, CipherResult, CipherType};

#[cfg(feature = "ring-aead-ciphers")]
use crate::crypto::ring::RingAeadCipher;
#[cfg(feature = "aes-pmac-siv")]
use crate::crypto::siv::MiscreantCipher;
#[cfg(feature = "sodium")]
use crate::crypto::sodium::SodiumAeadCipher;

use bytes::{Bytes, BytesMut};
use cfg_if::cfg_if;
use hkdf::Hkdf;
use sha1::Sha1;

/// Encryptor API for AEAD ciphers
pub trait AeadEncryptor {
    /// Encrypt `input` to `output` with `tag`. `output.len()` should equals to `input.len() + tag.len()`.
    /// ```plain
    /// +----------------------------------------+-----------------------+
    /// | ENCRYPTED TEXT (length = input.len())  | TAG                   |
    /// +----------------------------------------+-----------------------+
    /// ```
    fn encrypt(&mut self, input: &[u8], output: &mut [u8]);
}

/// Decryptor API for AEAD ciphers
pub trait AeadDecryptor {
    /// Decrypt `input` to `output` with `tag`. `output.len()` should equals to `input.len() - tag.len()`.
    /// ```plain
    /// +----------------------------------------+-----------------------+
    /// | ENCRYPTED TEXT (length = output.len()) | TAG                   |
    /// +----------------------------------------+-----------------------+
    /// ```
    fn decrypt(&mut self, input: &[u8], output: &mut [u8]) -> CipherResult<()>;
}

/// Variant `AeadDecryptor`
pub type BoxAeadDecryptor = Box<dyn AeadDecryptor + Send + 'static>;

/// Variant `AeadEncryptor`
pub type BoxAeadEncryptor = Box<dyn AeadEncryptor + Send + 'static>;

/// Generate a specific AEAD cipher encryptor
pub fn new_aead_encryptor(t: CipherType, key: &[u8], nonce: &[u8]) -> BoxAeadEncryptor {
    assert!(t.category() == CipherCategory::Aead);

    match t {
        #[cfg(feature = "ring-aead-ciphers")]
        CipherType::Aes128Gcm | CipherType::Aes256Gcm | CipherType::ChaCha20IetfPoly1305 => {
            Box::new(RingAeadCipher::new(t, key, nonce, true))
        }

        #[cfg(feature = "sodium")]
        CipherType::XChaCha20IetfPoly1305 => Box::new(SodiumAeadCipher::new(t, key, nonce)),

        #[cfg(feature = "aes-pmac-siv")]
        CipherType::Aes128PmacSiv | CipherType::Aes256PmacSiv => Box::new(MiscreantCipher::new(t, key, nonce)),

        _ => unreachable!(),
    }
}

/// Generate a specific AEAD cipher decryptor
pub fn new_aead_decryptor(t: CipherType, key: &[u8], nonce: &[u8]) -> BoxAeadDecryptor {
    assert!(t.category() == CipherCategory::Aead);

    match t {
        #[cfg(feature = "ring-aead-ciphers")]
        CipherType::Aes128Gcm | CipherType::Aes256Gcm | CipherType::ChaCha20IetfPoly1305 => {
            Box::new(RingAeadCipher::new(t, key, nonce, false))
        }

        #[cfg(feature = "sodium")]
        CipherType::XChaCha20IetfPoly1305 => Box::new(SodiumAeadCipher::new(t, key, nonce)),

        #[cfg(feature = "aes-pmac-siv")]
        CipherType::Aes128PmacSiv | CipherType::Aes256PmacSiv => Box::new(MiscreantCipher::new(t, key, nonce)),

        _ => unreachable!(),
    }
}

const SUBKEY_INFO: &[u8] = b"ss-subkey";

/// Make Session key
///
/// ## Session key (SIP007)
///
/// AEAD ciphers require a per-session subkey derived from the pre-shared master key using HKDF, and use the subkey
/// to encrypt/decrypt. Essentially it means we are moving from (M+N)-bit (PSK, nonce) pair to
/// (M+N)-bit (HKDF(PSK, salt), nonce) pair. Because HKDF is a PRF, the new construction significantly expands the
/// amount of randomness (from N to at least M where M is much greater than N), thus correcting the previously
/// mentioned design flaw.
///
/// Assuming we already have a user-supplied pre-shared master key PSK.
///
/// Function HKDF_SHA1 is a HKDF constructed using SHA1 hash. Its signature is
///
/// ```plain
/// HKDF_SHA1(secret_key, salt, info)
/// ```
///
/// The "info" string argument allows us to bind the derived subkey to a specific application context.
///
/// For AEAD ciphers, the encryption scheme is:
///
/// 1. Pick a random R-bit salt (R = max(128, len(SK)))
/// 2. Derive subkey SK = HKDF_SHA1(PSK, salt, "ss-subkey")
/// 3. Send salt
/// 4. For each chunk, encrypt and authenticate payload using SK with a counting nonce
///    (starting from 0 and increment by 1 after each use)
/// 5. Send encrypted chunk
pub fn make_skey(t: CipherType, key: &[u8], salt: &[u8]) -> Bytes {
    assert!(t.category() == CipherCategory::Aead);

    let hkdf = Hkdf::<Sha1>::new(Some(salt), key);

    let mut skey = BytesMut::with_capacity(key.len());
    unsafe {
        skey.set_len(key.len());
    }

    hkdf.expand(SUBKEY_INFO, &mut skey).unwrap();

    skey.freeze()
}

cfg_if! {
    if #[cfg(feature = "sodium")] {
        /// Increase nonce by 1
        ///
        /// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk (Little-Endian)
        pub fn increase_nonce(nonce: &mut [u8]) {
            use libsodium_sys::sodium_increment;

            unsafe {
                sodium_increment(nonce.as_mut_ptr(), nonce.len());
            }
        }
    } else {
        if #[cfg(target_endian = "little")] {
            if #[cfg(target_pointer_width = "32")] {
                /// Increase nonce by 1
                ///
                /// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk (Little-Endian)
                #[inline]
                pub fn increase_nonce(nonce: &mut [u8]) {
                    increase_nonce_32_le(nonce);
                }

            } else if #[cfg(target_pointer_width = "64")] {
                /// Increase nonce by 1
                ///
                /// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk (Little-Endian)
                #[inline]
                pub fn increase_nonce(nonce: &mut [u8]) {
                    increase_nonce_64_le(nonce);
                }
            } else {
                /// Increase nonce by 1
                ///
                /// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk (Little-Endian)
                #[inline]
                pub fn increase_nonce(nonce: &mut [u8]) {
                    increase_nonce_classic_le(nonce);
                }
            }

        } else {
            /// Increase nonce by 1
            ///
            /// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk (Little-Endian)
            #[inline]
            pub fn increase_nonce(nonce: &mut [u8]) {
                // Since current arthitecture is BE, integer add operations must be called after bswap operation.
                // bswap operation in most platforms are combinations of many instructions.
                increase_nonce_classic_le(nonce);
            }
        }
    }
}

#[allow(dead_code)]
fn increase_nonce_classic_le(nonce: &mut [u8]) {
    for i in nonce {
        if u8::MAX == *i {
            *i = 0;
        } else {
            *i += 1;
            return;
        }
    }
}

#[allow(dead_code)]
fn increase_nonce_32_le(nonce: &mut [u8]) {
    let data = nonce.as_mut_ptr();

    let mut i = 0;
    while i + 4 < nonce.len() {
        let v = unsafe { &mut *(data.offset(i as isize) as *mut u32) };
        if *v == u32::MAX {
            *v = 0;
        } else {
            *v += 1;
            return;
        }
        i += 4;
    }

    if i < nonce.len() {
        // unlikely (missing intrinsic)

        increase_nonce_classic_le(&mut nonce[i..]);
    }
}

#[allow(dead_code)]
fn increase_nonce_64_le(nonce: &mut [u8]) {
    let data = nonce.as_mut_ptr();

    let mut i = 0;
    while i + 8 < nonce.len() {
        let v = unsafe { &mut *(data.offset(i as isize) as *mut u64) };
        if *v == u64::MAX {
            *v = 0;
        } else {
            *v += 1;
            return;
        }
        i += 8;
    }

    if i < nonce.len() {
        // unlinkely (missing intrinsic)

        if i + 4 < nonce.len() {
            increase_nonce_32_le(&mut nonce[i..]);
        } else {
            increase_nonce_classic_le(&mut nonce[i..]);
        }
    }
}

#[cfg(test)]
mod test {
    #[test]
    fn increase_nonce_classic_le() {
        let mut v1 = [0, 0, 0, 0];
        super::increase_nonce_classic_le(&mut v1);
        assert_eq!(v1, [1, 0, 0, 0]);

        let mut v2 = [0xff, 0, 0, 0];
        super::increase_nonce_classic_le(&mut v2);
        assert_eq!(v2, [0, 1, 0, 0]);

        let mut v3 = [0xff, 0xff, 0xff, 0xff];
        super::increase_nonce_classic_le(&mut v3);
        assert_eq!(v3, [0, 0, 0, 0]);
    }

    #[test]
    #[cfg(target_endian = "little")]
    fn increase_nonce_32_le() {
        let mut v1 = [0, 0, 0, 0];
        super::increase_nonce_32_le(&mut v1);
        assert_eq!(v1, [1, 0, 0, 0]);

        let mut v2 = [0xff, 0, 0, 0];
        super::increase_nonce_32_le(&mut v2);
        assert_eq!(v2, [0, 1, 0, 0]);

        let mut v3 = [0xff, 0xff, 0xff, 0xff];
        super::increase_nonce_32_le(&mut v3);
        assert_eq!(v3, [0, 0, 0, 0]);

        let mut v4 = [0, 0, 0, 0, 0, 0];
        super::increase_nonce_32_le(&mut v4);
        assert_eq!(v4, [1, 0, 0, 0, 0, 0]);

        let mut v5 = [0xff, 0xff, 0xff, 0xff, 0, 0];
        super::increase_nonce_32_le(&mut v5);
        assert_eq!(v5, [0, 0, 0, 0, 1, 0]);
    }

    #[test]
    #[cfg(target_endian = "little")]
    fn increase_nonce_64_le() {
        let mut v1 = [0, 0, 0, 0, 0, 0, 0, 0];
        super::increase_nonce_64_le(&mut v1);
        assert_eq!(v1, [1, 0, 0, 0, 0, 0, 0, 0]);

        let mut v2 = [0xff, 0, 0, 0, 0, 0, 0, 0];
        super::increase_nonce_64_le(&mut v2);
        assert_eq!(v2, [0, 1, 0, 0, 0, 0, 0, 0]);

        let mut v3 = [0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff];
        super::increase_nonce_64_le(&mut v3);
        assert_eq!(v3, [0, 0, 0, 0, 0, 0, 0, 0]);

        let mut v4 = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        super::increase_nonce_64_le(&mut v4);
        assert_eq!(v4, [1, 0, 0, 0, 0, 0, 0, 0, 0, 0]);

        let mut v5 = [0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0, 0];
        super::increase_nonce_64_le(&mut v5);
        assert_eq!(v5, [0, 0, 0, 0, 0, 0, 0, 0, 1, 0]);

        let mut v6 = [0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0, 0, 0];
        super::increase_nonce_64_le(&mut v6);
        assert_eq!(v6, [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0]);
    }
}
<commit_msg>cannot avoid judging of multiple of 4, just call increase_nonce_32_le<commit_after>//! Aead Ciphers

use crate::crypto::cipher::{CipherCategory, CipherResult, CipherType};

#[cfg(feature = "ring-aead-ciphers")]
use crate::crypto::ring::RingAeadCipher;
#[cfg(feature = "aes-pmac-siv")]
use crate::crypto::siv::MiscreantCipher;
#[cfg(feature = "sodium")]
use crate::crypto::sodium::SodiumAeadCipher;

use bytes::{Bytes, BytesMut};
use cfg_if::cfg_if;
use hkdf::Hkdf;
use sha1::Sha1;

/// Encryptor API for AEAD ciphers
pub trait AeadEncryptor {
    /// Encrypt `input` to `output` with `tag`. `output.len()` should equals to `input.len() + tag.len()`.
    /// ```plain
    /// +----------------------------------------+-----------------------+
    /// | ENCRYPTED TEXT (length = input.len())  | TAG                   |
    /// +----------------------------------------+-----------------------+
    /// ```
    fn encrypt(&mut self, input: &[u8], output: &mut [u8]);
}

/// Decryptor API for AEAD ciphers
pub trait AeadDecryptor {
    /// Decrypt `input` to `output` with `tag`. `output.len()` should equals to `input.len() - tag.len()`.
    /// ```plain
    /// +----------------------------------------+-----------------------+
    /// | ENCRYPTED TEXT (length = output.len()) | TAG                   |
    /// +----------------------------------------+-----------------------+
    /// ```
    fn decrypt(&mut self, input: &[u8], output: &mut [u8]) -> CipherResult<()>;
}

/// Variant `AeadDecryptor`
pub type BoxAeadDecryptor = Box<dyn AeadDecryptor + Send + 'static>;

/// Variant `AeadEncryptor`
pub type BoxAeadEncryptor = Box<dyn AeadEncryptor + Send + 'static>;

/// Generate a specific AEAD cipher encryptor
pub fn new_aead_encryptor(t: CipherType, key: &[u8], nonce: &[u8]) -> BoxAeadEncryptor {
    assert!(t.category() == CipherCategory::Aead);

    match t {
        #[cfg(feature = "ring-aead-ciphers")]
        CipherType::Aes128Gcm | CipherType::Aes256Gcm | CipherType::ChaCha20IetfPoly1305 => {
            Box::new(RingAeadCipher::new(t, key, nonce, true))
        }

        #[cfg(feature = "sodium")]
        CipherType::XChaCha20IetfPoly1305 => Box::new(SodiumAeadCipher::new(t, key, nonce)),

        #[cfg(feature = "aes-pmac-siv")]
        CipherType::Aes128PmacSiv | CipherType::Aes256PmacSiv => Box::new(MiscreantCipher::new(t, key, nonce)),

        _ => unreachable!(),
    }
}

/// Generate a specific AEAD cipher decryptor
pub fn new_aead_decryptor(t: CipherType, key: &[u8], nonce: &[u8]) -> BoxAeadDecryptor {
    assert!(t.category() == CipherCategory::Aead);

    match t {
        #[cfg(feature = "ring-aead-ciphers")]
        CipherType::Aes128Gcm | CipherType::Aes256Gcm | CipherType::ChaCha20IetfPoly1305 => {
            Box::new(RingAeadCipher::new(t, key, nonce, false))
        }

        #[cfg(feature = "sodium")]
        CipherType::XChaCha20IetfPoly1305 => Box::new(SodiumAeadCipher::new(t, key, nonce)),

        #[cfg(feature = "aes-pmac-siv")]
        CipherType::Aes128PmacSiv | CipherType::Aes256PmacSiv => Box::new(MiscreantCipher::new(t, key, nonce)),

        _ => unreachable!(),
    }
}

const SUBKEY_INFO: &[u8] = b"ss-subkey";

/// Make Session key
///
/// ## Session key (SIP007)
///
/// AEAD ciphers require a per-session subkey derived from the pre-shared master key using HKDF, and use the subkey
/// to encrypt/decrypt. Essentially it means we are moving from (M+N)-bit (PSK, nonce) pair to
/// (M+N)-bit (HKDF(PSK, salt), nonce) pair. Because HKDF is a PRF, the new construction significantly expands the
/// amount of randomness (from N to at least M where M is much greater than N), thus correcting the previously
/// mentioned design flaw.
///
/// Assuming we already have a user-supplied pre-shared master key PSK.
///
/// Function HKDF_SHA1 is a HKDF constructed using SHA1 hash. Its signature is
///
/// ```plain
/// HKDF_SHA1(secret_key, salt, info)
/// ```
///
/// The "info" string argument allows us to bind the derived subkey to a specific application context.
///
/// For AEAD ciphers, the encryption scheme is:
///
/// 1. Pick a random R-bit salt (R = max(128, len(SK)))
/// 2. Derive subkey SK = HKDF_SHA1(PSK, salt, "ss-subkey")
/// 3. Send salt
/// 4. For each chunk, encrypt and authenticate payload using SK with a counting nonce
///    (starting from 0 and increment by 1 after each use)
/// 5. Send encrypted chunk
pub fn make_skey(t: CipherType, key: &[u8], salt: &[u8]) -> Bytes {
    assert!(t.category() == CipherCategory::Aead);

    let hkdf = Hkdf::<Sha1>::new(Some(salt), key);

    let mut skey = BytesMut::with_capacity(key.len());
    unsafe {
        skey.set_len(key.len());
    }

    hkdf.expand(SUBKEY_INFO, &mut skey).unwrap();

    skey.freeze()
}

cfg_if! {
    if #[cfg(feature = "sodium")] {
        /// Increase nonce by 1
        ///
        /// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk (Little-Endian)
        pub fn increase_nonce(nonce: &mut [u8]) {
            use libsodium_sys::sodium_increment;

            unsafe {
                sodium_increment(nonce.as_mut_ptr(), nonce.len());
            }
        }
    } else {
        if #[cfg(target_endian = "little")] {
            if #[cfg(target_pointer_width = "32")] {
                /// Increase nonce by 1
                ///
                /// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk (Little-Endian)
                #[inline]
                pub fn increase_nonce(nonce: &mut [u8]) {
                    increase_nonce_32_le(nonce);
                }

            } else if #[cfg(target_pointer_width = "64")] {
                /// Increase nonce by 1
                ///
                /// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk (Little-Endian)
                #[inline]
                pub fn increase_nonce(nonce: &mut [u8]) {
                    increase_nonce_64_le(nonce);
                }
            } else {
                /// Increase nonce by 1
                ///
                /// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk (Little-Endian)
                #[inline]
                pub fn increase_nonce(nonce: &mut [u8]) {
                    increase_nonce_classic_le(nonce);
                }
            }

        } else {
            /// Increase nonce by 1
            ///
            /// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk (Little-Endian)
            #[inline]
            pub fn increase_nonce(nonce: &mut [u8]) {
                // Since current arthitecture is BE, integer add operations must be called after bswap operation.
                // bswap operation in most platforms are combinations of many instructions.
                increase_nonce_classic_le(nonce);
            }
        }
    }
}

#[allow(dead_code)]
fn increase_nonce_classic_le(nonce: &mut [u8]) {
    for i in nonce {
        if u8::MAX == *i {
            *i = 0;
        } else {
            *i += 1;
            return;
        }
    }
}

#[allow(dead_code)]
fn increase_nonce_32_le(nonce: &mut [u8]) {
    let data = nonce.as_mut_ptr();

    let mut i = 0;
    while i + 4 < nonce.len() {
        let v = unsafe { &mut *(data.offset(i as isize) as *mut u32) };
        if *v == u32::MAX {
            *v = 0;
        } else {
            *v += 1;
            return;
        }
        i += 4;
    }

    if i < nonce.len() {
        // unlikely (missing intrinsic)

        increase_nonce_classic_le(&mut nonce[i..]);
    }
}

#[allow(dead_code)]
fn increase_nonce_64_le(nonce: &mut [u8]) {
    let data = nonce.as_mut_ptr();

    let mut i = 0;
    while i + 8 < nonce.len() {
        let v = unsafe { &mut *(data.offset(i as isize) as *mut u64) };
        if *v == u64::MAX {
            *v = 0;
        } else {
            *v += 1;
            return;
        }
        i += 8;
    }

    if i < nonce.len() {
        // unlinkely (missing intrinsic)

        increase_nonce_32_le(&mut nonce[i..]);
    }
}

#[cfg(test)]
mod test {
    #[test]
    fn increase_nonce_classic_le() {
        let mut v1 = [0, 0, 0, 0];
        super::increase_nonce_classic_le(&mut v1);
        assert_eq!(v1, [1, 0, 0, 0]);

        let mut v2 = [0xff, 0, 0, 0];
        super::increase_nonce_classic_le(&mut v2);
        assert_eq!(v2, [0, 1, 0, 0]);

        let mut v3 = [0xff, 0xff, 0xff, 0xff];
        super::increase_nonce_classic_le(&mut v3);
        assert_eq!(v3, [0, 0, 0, 0]);
    }

    #[test]
    #[cfg(target_endian = "little")]
    fn increase_nonce_32_le() {
        let mut v1 = [0, 0, 0, 0];
        super::increase_nonce_32_le(&mut v1);
        assert_eq!(v1, [1, 0, 0, 0]);

        let mut v2 = [0xff, 0, 0, 0];
        super::increase_nonce_32_le(&mut v2);
        assert_eq!(v2, [0, 1, 0, 0]);

        let mut v3 = [0xff, 0xff, 0xff, 0xff];
        super::increase_nonce_32_le(&mut v3);
        assert_eq!(v3, [0, 0, 0, 0]);

        let mut v4 = [0, 0, 0, 0, 0, 0];
        super::increase_nonce_32_le(&mut v4);
        assert_eq!(v4, [1, 0, 0, 0, 0, 0]);

        let mut v5 = [0xff, 0xff, 0xff, 0xff, 0, 0];
        super::increase_nonce_32_le(&mut v5);
        assert_eq!(v5, [0, 0, 0, 0, 1, 0]);
    }

    #[test]
    #[cfg(target_endian = "little")]
    fn increase_nonce_64_le() {
        let mut v1 = [0, 0, 0, 0, 0, 0, 0, 0];
        super::increase_nonce_64_le(&mut v1);
        assert_eq!(v1, [1, 0, 0, 0, 0, 0, 0, 0]);

        let mut v2 = [0xff, 0, 0, 0, 0, 0, 0, 0];
        super::increase_nonce_64_le(&mut v2);
        assert_eq!(v2, [0, 1, 0, 0, 0, 0, 0, 0]);

        let mut v3 = [0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff];
        super::increase_nonce_64_le(&mut v3);
        assert_eq!(v3, [0, 0, 0, 0, 0, 0, 0, 0]);

        let mut v4 = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        super::increase_nonce_64_le(&mut v4);
        assert_eq!(v4, [1, 0, 0, 0, 0, 0, 0, 0, 0, 0]);

        let mut v5 = [0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0, 0];
        super::increase_nonce_64_le(&mut v5);
        assert_eq!(v5, [0, 0, 0, 0, 0, 0, 0, 0, 1, 0]);

        let mut v6 = [0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0, 0, 0];
        super::increase_nonce_64_le(&mut v6);
        assert_eq!(v6, [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0]);
    }
}
<|endoftext|>


<commit_before>use objects::message::Message;
use objects::user::User;

#[derive(Clone, Deserialize, Debug)]
pub struct CallBackQuery {
	id: String,
	from: User,
	message: Option<Message>,
	inline_message_id: Option<String>,
	chat_instance: String,
	data: Option<String>,
	game_short_name: Option<String>,
}<commit_msg>file apeared to contain tabs instead op spaces, converted it to spaces now<commit_after>use objects::message::Message;
use objects::user::User;

#[derive(Clone, Deserialize, Debug)]
pub struct CallBackQuery {
    id: String,
    from: User,
    message: Option<Message>,
    inline_message_id: Option<String>,
    chat_instance: String,
    data: Option<String>,
    game_short_name: Option<String>,
}<|endoftext|>


<commit_before>// The MIT License (MIT)

// Copyright (c) 2015 Y. T. Chung <zonyitoo@gmail.com>

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

//! Decoder

mod error;
mod serde;

pub use self::error::{DecoderError, DecoderResult};
pub use self::serde::Decoder;

use std::io::Read;

use byteorder::{LittleEndian, ReadBytesExt};
use chrono::offset::{LocalResult, TimeZone};
use chrono::Utc;

use bson::{Array, Bson, Document};
use oid;
use spec::{self, BinarySubtype};

use serde::de::Deserialize;

fn read_string<R: Read + ?Sized>(reader: &mut R, utf8_lossy: bool) -> DecoderResult<String> {
    let len = reader.read_i32::<LittleEndian>()?;

    let s = if utf8_lossy {
        let mut buf = Vec::with_capacity(len as usize - 1);
        reader.take(len as u64 - 1).read_to_end(&mut buf)?;
        String::from_utf8_lossy(&buf).to_string()
    } else {
        let mut s = String::with_capacity(len as usize - 1);
        reader.take(len as u64 - 1).read_to_string(&mut s)?;
        s
    };
    reader.read_u8()?; // The last 0x00

    Ok(s)
}

fn read_cstring<R: Read + ?Sized>(reader: &mut R) -> DecoderResult<String> {
    let mut v = Vec::new();

    loop {
        let c = reader.read_u8()?;
        if c == 0 {
            break;
        }
        v.push(c);
    }

    Ok(String::from_utf8(v)?)
}

#[inline]
fn read_i32<R: Read + ?Sized>(reader: &mut R) -> DecoderResult<i32> {
    reader.read_i32::<LittleEndian>().map_err(From::from)
}

#[inline]
fn read_i64<R: Read + ?Sized>(reader: &mut R) -> DecoderResult<i64> {
    reader.read_i64::<LittleEndian>().map_err(From::from)
}

/// Attempt to decode a `Document` from a byte stream.
pub fn decode_document<R: Read + ?Sized>(reader: &mut R) -> DecoderResult<Document> {
    let mut doc = Document::new();

    // disregard the length: using Read::take causes infinite type recursion
    read_i32(reader)?;

    loop {
        let tag = reader.read_u8()?;

        if tag == 0 {
            break;
        }

        let key = read_cstring(reader)?;
        let val = decode_bson(reader, tag, false)?;

        doc.insert(key, val);
    }

    Ok(doc)
}

/// Attempt to decode a `Document` that may contain invalid UTF-8 strings from a byte stream.
pub fn decode_document_utf8_lossy<R: Read + ?Sized>(reader: &mut R) -> DecoderResult<Document> {
    let mut doc = Document::new();

    // disregard the length: using Read::take causes infinite type recursion
    read_i32(reader)?;

    loop {
        let tag = reader.read_u8()?;

        if tag == 0 {
            break;
        }

        let key = read_cstring(reader)?;
        let val = decode_bson(reader, tag, true)?;

        doc.insert(key, val);
    }

    Ok(doc)
}

fn decode_array<R: Read + ?Sized>(reader: &mut R, utf8_lossy: bool) -> DecoderResult<Array> {
    let mut arr = Array::new();

    // disregard the length: using Read::take causes infinite type recursion
    read_i32(reader)?;

    loop {
        let tag = reader.read_u8()?;
        if tag == 0 {
            break;
        }

        // check that the key is as expected
        let key = read_cstring(reader)?;
        match key.parse::<usize>() {
            Err(..) => return Err(DecoderError::InvalidArrayKey(arr.len(), key)),
            Ok(idx) => {
                if idx != arr.len() {
                    return Err(DecoderError::InvalidArrayKey(arr.len(), key));
                }
            }
        }

        let val = decode_bson(reader, tag, utf8_lossy)?;
        arr.push(val)
    }

    Ok(arr)
}

fn decode_bson<R: Read + ?Sized>(reader: &mut R, tag: u8, utf8_lossy: bool) -> DecoderResult<Bson> {
    use spec::ElementType::*;
    match spec::ElementType::from(tag) {
        Some(FloatingPoint) => Ok(Bson::FloatingPoint(reader.read_f64::<LittleEndian>()?)),
        Some(Utf8String) => read_string(reader, utf8_lossy).map(Bson::String),
        Some(EmbeddedDocument) => decode_document(reader).map(Bson::Document),
        Some(Array) => decode_array(reader, utf8_lossy).map(Bson::Array),
        Some(Binary) => {
            let len = read_i32(reader)?;
            let subtype = BinarySubtype::from(reader.read_u8()?);
            let mut data = Vec::with_capacity(len as usize);
            reader.take(len as u64).read_to_end(&mut data)?;
            Ok(Bson::Binary(subtype, data))
        }
        Some(ObjectId) => {
            let mut objid = [0; 12];
            for x in &mut objid {
                *x = reader.read_u8()?;
            }
            Ok(Bson::ObjectId(oid::ObjectId::with_bytes(objid)))
        }
        Some(Boolean) => Ok(Bson::Boolean(reader.read_u8()? != 0)),
        Some(NullValue) => Ok(Bson::Null),
        Some(RegularExpression) => {
            let pat = read_cstring(reader)?;
            let opt = read_cstring(reader)?;
            Ok(Bson::RegExp(pat, opt))
        }
        Some(JavaScriptCode) => read_string(reader, utf8_lossy).map(Bson::JavaScriptCode),
        Some(JavaScriptCodeWithScope) => {
            // disregard the length:
            //     using Read::take causes infinite type recursion
            read_i32(reader)?;

            let code = read_string(reader, utf8_lossy)?;
            let scope = decode_document(reader)?;
            Ok(Bson::JavaScriptCodeWithScope(code, scope))
        }
        Some(Integer32Bit) => read_i32(reader).map(Bson::I32),
        Some(Integer64Bit) => read_i64(reader).map(Bson::I64),
        Some(TimeStamp) => read_i64(reader).map(Bson::TimeStamp),
        Some(UtcDatetime) => {
            let time = read_i64(reader)?;

            match Utc.timestamp_opt(time / 1000, (time % 1000) as u32 * 1000000) {
                LocalResult::None => Err(DecoderError::InvalidTimestamp(time)),
                LocalResult::Ambiguous(..) => Err(DecoderError::AmbiguousTimestamp(time)),
                LocalResult::Single(t) => Ok(Bson::UtcDatetime(t)),
            }
        }
        Some(Symbol) => read_string(reader, utf8_lossy).map(Bson::Symbol),
        Some(Undefined) | Some(DbPointer) | Some(MaxKey) | Some(MinKey) | None => {
            Err(DecoderError::UnrecognizedElementType(tag))
        }
    }
}

/// Decode a BSON `Value` into a `T` Deserializable.
pub fn from_bson<'de, T>(bson: Bson) -> DecoderResult<T>
    where T: Deserialize<'de>
{
    let de = Decoder::new(bson);
    Deserialize::deserialize(de)
}
<commit_msg>[#64] Do not panic if length of UTF-8 string is invalid<commit_after>// The MIT License (MIT)

// Copyright (c) 2015 Y. T. Chung <zonyitoo@gmail.com>

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

//! Decoder

mod error;
mod serde;

pub use self::error::{DecoderError, DecoderResult};
pub use self::serde::Decoder;

use std::io::Read;

use byteorder::{LittleEndian, ReadBytesExt};
use chrono::offset::{LocalResult, TimeZone};
use chrono::Utc;

use bson::{Array, Bson, Document};
use oid;
use spec::{self, BinarySubtype};

use serde::de::Deserialize;

fn read_string<R: Read + ?Sized>(reader: &mut R, utf8_lossy: bool) -> DecoderResult<String> {
    let len = reader.read_i32::<LittleEndian>()?;

    // UTF-8 String must have at least 1 byte (the last 0x00).
    if len < 1 {
        return Err(DecoderError::InvalidLength(len as usize, format!("invalid length {} for UTF-8 string", len)));
    }

    let s = if utf8_lossy {
        let mut buf = Vec::with_capacity(len as usize - 1);
        reader.take(len as u64 - 1).read_to_end(&mut buf)?;
        String::from_utf8_lossy(&buf).to_string()
    } else {
        let mut s = String::with_capacity(len as usize - 1);
        reader.take(len as u64 - 1).read_to_string(&mut s)?;
        s
    };
    reader.read_u8()?; // The last 0x00

    Ok(s)
}

fn read_cstring<R: Read + ?Sized>(reader: &mut R) -> DecoderResult<String> {
    let mut v = Vec::new();

    loop {
        let c = reader.read_u8()?;
        if c == 0 {
            break;
        }
        v.push(c);
    }

    Ok(String::from_utf8(v)?)
}

#[inline]
fn read_i32<R: Read + ?Sized>(reader: &mut R) -> DecoderResult<i32> {
    reader.read_i32::<LittleEndian>().map_err(From::from)
}

#[inline]
fn read_i64<R: Read + ?Sized>(reader: &mut R) -> DecoderResult<i64> {
    reader.read_i64::<LittleEndian>().map_err(From::from)
}

/// Attempt to decode a `Document` from a byte stream.
pub fn decode_document<R: Read + ?Sized>(reader: &mut R) -> DecoderResult<Document> {
    let mut doc = Document::new();

    // disregard the length: using Read::take causes infinite type recursion
    read_i32(reader)?;

    loop {
        let tag = reader.read_u8()?;

        if tag == 0 {
            break;
        }

        let key = read_cstring(reader)?;
        let val = decode_bson(reader, tag, false)?;

        doc.insert(key, val);
    }

    Ok(doc)
}

/// Attempt to decode a `Document` that may contain invalid UTF-8 strings from a byte stream.
pub fn decode_document_utf8_lossy<R: Read + ?Sized>(reader: &mut R) -> DecoderResult<Document> {
    let mut doc = Document::new();

    // disregard the length: using Read::take causes infinite type recursion
    read_i32(reader)?;

    loop {
        let tag = reader.read_u8()?;

        if tag == 0 {
            break;
        }

        let key = read_cstring(reader)?;
        let val = decode_bson(reader, tag, true)?;

        doc.insert(key, val);
    }

    Ok(doc)
}

fn decode_array<R: Read + ?Sized>(reader: &mut R, utf8_lossy: bool) -> DecoderResult<Array> {
    let mut arr = Array::new();

    // disregard the length: using Read::take causes infinite type recursion
    read_i32(reader)?;

    loop {
        let tag = reader.read_u8()?;
        if tag == 0 {
            break;
        }

        // check that the key is as expected
        let key = read_cstring(reader)?;
        match key.parse::<usize>() {
            Err(..) => return Err(DecoderError::InvalidArrayKey(arr.len(), key)),
            Ok(idx) => {
                if idx != arr.len() {
                    return Err(DecoderError::InvalidArrayKey(arr.len(), key));
                }
            }
        }

        let val = decode_bson(reader, tag, utf8_lossy)?;
        arr.push(val)
    }

    Ok(arr)
}

fn decode_bson<R: Read + ?Sized>(reader: &mut R, tag: u8, utf8_lossy: bool) -> DecoderResult<Bson> {
    use spec::ElementType::*;
    match spec::ElementType::from(tag) {
        Some(FloatingPoint) => Ok(Bson::FloatingPoint(reader.read_f64::<LittleEndian>()?)),
        Some(Utf8String) => read_string(reader, utf8_lossy).map(Bson::String),
        Some(EmbeddedDocument) => decode_document(reader).map(Bson::Document),
        Some(Array) => decode_array(reader, utf8_lossy).map(Bson::Array),
        Some(Binary) => {
            let len = read_i32(reader)?;
            let subtype = BinarySubtype::from(reader.read_u8()?);
            let mut data = Vec::with_capacity(len as usize);
            reader.take(len as u64).read_to_end(&mut data)?;
            Ok(Bson::Binary(subtype, data))
        }
        Some(ObjectId) => {
            let mut objid = [0; 12];
            for x in &mut objid {
                *x = reader.read_u8()?;
            }
            Ok(Bson::ObjectId(oid::ObjectId::with_bytes(objid)))
        }
        Some(Boolean) => Ok(Bson::Boolean(reader.read_u8()? != 0)),
        Some(NullValue) => Ok(Bson::Null),
        Some(RegularExpression) => {
            let pat = read_cstring(reader)?;
            let opt = read_cstring(reader)?;
            Ok(Bson::RegExp(pat, opt))
        }
        Some(JavaScriptCode) => read_string(reader, utf8_lossy).map(Bson::JavaScriptCode),
        Some(JavaScriptCodeWithScope) => {
            // disregard the length:
            //     using Read::take causes infinite type recursion
            read_i32(reader)?;

            let code = read_string(reader, utf8_lossy)?;
            let scope = decode_document(reader)?;
            Ok(Bson::JavaScriptCodeWithScope(code, scope))
        }
        Some(Integer32Bit) => read_i32(reader).map(Bson::I32),
        Some(Integer64Bit) => read_i64(reader).map(Bson::I64),
        Some(TimeStamp) => read_i64(reader).map(Bson::TimeStamp),
        Some(UtcDatetime) => {
            let time = read_i64(reader)?;

            match Utc.timestamp_opt(time / 1000, (time % 1000) as u32 * 1000000) {
                LocalResult::None => Err(DecoderError::InvalidTimestamp(time)),
                LocalResult::Ambiguous(..) => Err(DecoderError::AmbiguousTimestamp(time)),
                LocalResult::Single(t) => Ok(Bson::UtcDatetime(t)),
            }
        }
        Some(Symbol) => read_string(reader, utf8_lossy).map(Bson::Symbol),
        Some(Undefined) | Some(DbPointer) | Some(MaxKey) | Some(MinKey) | None => {
            Err(DecoderError::UnrecognizedElementType(tag))
        }
    }
}

/// Decode a BSON `Value` into a `T` Deserializable.
pub fn from_bson<'de, T>(bson: Bson) -> DecoderResult<T>
    where T: Deserialize<'de>
{
    let de = Decoder::new(bson);
    Deserialize::deserialize(de)
}
<|endoftext|>


<commit_before>use crate::{
    analysis::{rust_type::RustType, trampoline_parameters::Transformation},
    env::Env,
    library,
    nameutil::is_gstring,
    traits::*,
};

pub trait TrampolineFromGlib {
    fn trampoline_from_glib(&self, env: &Env, need_downcast: bool, nullable: bool) -> String;
}

impl TrampolineFromGlib for Transformation {
    fn trampoline_from_glib(&self, env: &Env, need_downcast: bool, nullable: bool) -> String {
        use crate::analysis::conversion_type::ConversionType::*;
        let need_type_name = need_downcast || is_need_type_name(env, self.typ);
        match self.conversion_type {
            Direct => self.name.clone(),
            Scalar | Option | Result { .. } => format!("from_glib({})", self.name),
            Borrow | Pointer => {
                let is_borrow = self.conversion_type == Borrow;
                let need_type_name = need_type_name || (is_borrow && nullable);
                let (mut left, mut right) = from_glib_xxx(self.transfer, is_borrow);
                let type_name = RustType::try_new(env, self.typ).into_string();
                if need_type_name {
                    if is_borrow && nullable {
                        left = format!("Option::<{}>::{}", type_name, left);
                    } else {
                        left = format!("{}::{}", type_name, left);
                    }
                }

                if !nullable {
                    left = format!("&{}", left);
                } else if nullable && is_borrow {
                    if is_gstring(&type_name) {
                        right = format!("{}.as_ref().as_deref()", right);
                    } else {
                        right = format!("{}.as_ref().as_ref()", right);
                    }
                } else if is_gstring(&type_name) {
                    right = format!("{}.as_deref()", right);
                } else {
                    right = format!("{}.as_ref()", right);
                }

                if need_downcast && is_borrow {
                    right = format!("{}.unsafe_cast_ref()", right);
                } else if need_downcast {
                    right = format!("{}.unsafe_cast()", right);
                }

                format!("{}{}{}", left, self.name, right)
            }
            Unknown => format!("/*Unknown conversion*/{}", self.name),
        }
    }
}

pub fn from_glib_xxx(transfer: library::Transfer, is_borrow: bool) -> (String, String) {
    use crate::library::Transfer::*;
    match transfer {
        None if is_borrow => ("from_glib_borrow(".into(), ")".into()),
        None => ("from_glib_none(".into(), ")".into()),
        Full => ("from_glib_full(".into(), ")".into()),
        Container => ("from_glib_container(".into(), ")".into()),
    }
}

fn is_need_type_name(env: &Env, type_id: library::TypeId) -> bool {
    if type_id.ns_id == library::INTERNAL_NAMESPACE {
        use crate::library::{Fundamental::*, Type::*};
        matches!(
            env.type_(type_id),
            Fundamental(Utf8) | Fundamental(Filename) | Fundamental(OsString)
        )
    } else {
        false
    }
}
<commit_msg>Remove extra & to fix clippy lint<commit_after>use crate::{
    analysis::{rust_type::RustType, trampoline_parameters::Transformation},
    env::Env,
    library,
    nameutil::is_gstring,
    traits::*,
};

pub trait TrampolineFromGlib {
    fn trampoline_from_glib(&self, env: &Env, need_downcast: bool, nullable: bool) -> String;
}

impl TrampolineFromGlib for Transformation {
    fn trampoline_from_glib(&self, env: &Env, need_downcast: bool, nullable: bool) -> String {
        use crate::analysis::conversion_type::ConversionType::*;
        let need_type_name = need_downcast || is_need_type_name(env, self.typ);
        match self.conversion_type {
            Direct => self.name.clone(),
            Scalar | Option | Result { .. } => format!("from_glib({})", self.name),
            Borrow | Pointer => {
                let is_borrow = self.conversion_type == Borrow;
                let need_type_name = need_type_name || (is_borrow && nullable);
                let (mut left, mut right) = from_glib_xxx(self.transfer, is_borrow);
                let type_name = RustType::try_new(env, self.typ).into_string();
                if need_type_name {
                    if is_borrow && nullable {
                        left = format!("Option::<{}>::{}", type_name, left);
                    } else {
                        left = format!("{}::{}", type_name, left);
                    }
                }

                if !nullable {
                    left = format!(
                        "{}{}",
                        if need_downcast && is_borrow { "" } else { "&" },
                        left
                    );
                } else if nullable && is_borrow {
                    if is_gstring(&type_name) {
                        right = format!("{}.as_ref().as_deref()", right);
                    } else {
                        right = format!("{}.as_ref().as_ref()", right);
                    }
                } else if is_gstring(&type_name) {
                    right = format!("{}.as_deref()", right);
                } else {
                    right = format!("{}.as_ref()", right);
                }

                if need_downcast && is_borrow {
                    right = format!("{}.unsafe_cast_ref()", right);
                } else if need_downcast {
                    right = format!("{}.unsafe_cast()", right);
                }

                format!("{}{}{}", left, self.name, right)
            }
            Unknown => format!("/*Unknown conversion*/{}", self.name),
        }
    }
}

pub fn from_glib_xxx(transfer: library::Transfer, is_borrow: bool) -> (String, String) {
    use crate::library::Transfer::*;
    match transfer {
        None if is_borrow => ("from_glib_borrow(".into(), ")".into()),
        None => ("from_glib_none(".into(), ")".into()),
        Full => ("from_glib_full(".into(), ")".into()),
        Container => ("from_glib_container(".into(), ")".into()),
    }
}

fn is_need_type_name(env: &Env, type_id: library::TypeId) -> bool {
    if type_id.ns_id == library::INTERNAL_NAMESPACE {
        use crate::library::{Fundamental::*, Type::*};
        matches!(
            env.type_(type_id),
            Fundamental(Utf8) | Fundamental(Filename) | Fundamental(OsString)
        )
    } else {
        false
    }
}
<|endoftext|>


<commit_before>//! TODO Fill in

use std::fmt::{self, Display, Formatter};
use std::default::Default;
use rlua::{self, Table, Lua, UserData, ToLua};
use super::signal;

#[derive(Clone, Debug)]
pub struct MouseState {
    // TODO Fill in
    dummy: i32
}

impl Default for MouseState {
    fn default() -> Self {
        MouseState {
            dummy: 0
        }
    }
}

impl Display for MouseState {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "Mouse: {:p}", self)
    }
}

impl UserData for MouseState {}

pub fn init(lua: &Lua) -> rlua::Result<()> {
    let mouse_table = lua.create_table();
    state_setup(lua, &mouse_table)?;
    meta_setup(lua, &mouse_table)?;
    method_setup(lua, &mouse_table)?;
    let globals = lua.globals();
    globals.set("mouse", mouse_table)
}

fn state_setup(lua: &Lua, mouse_table: &Table) -> rlua::Result<()> {
    mouse_table.set("__data", MouseState::default().to_lua(lua)?)
}

fn meta_setup(lua: &Lua, mouse_table: &Table) -> rlua::Result<()> {
    let meta_table = lua.create_table();
    mouse_table.set("__tostring", lua.create_function(|_, val: Table| {
        Ok(format!("{}", val.get::<_, MouseState>("__data")?))
    }))?;
    mouse_table.set_metatable(Some(meta_table));
    Ok(())
}

fn method_setup(lua: &Lua, mouse_table: &Table) -> rlua::Result<()> {
    mouse_table.set("connect_signal", lua.create_function(signal::global_connect_signal))
}
<commit_msg>Added mouse.coords<commit_after>//! TODO Fill in

use std::fmt::{self, Display, Formatter};
use std::default::Default;
use rlua::{self, Table, Lua, UserData, ToLua};
use rustwlc::{Point, input};

#[derive(Clone, Debug)]
pub struct MouseState {
    // TODO Fill in
    dummy: i32
}

impl Default for MouseState {
    fn default() -> Self {
        MouseState {
            dummy: 0
        }
    }
}

impl Display for MouseState {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "Mouse: {:p}", self)
    }
}

impl UserData for MouseState {}

pub fn init(lua: &Lua) -> rlua::Result<()> {
    let mouse_table = lua.create_table();
    state_setup(lua, &mouse_table)?;
    meta_setup(lua, &mouse_table)?;
    method_setup(lua, &mouse_table)?;
    let globals = lua.globals();
    globals.set("mouse", mouse_table)
}

fn state_setup(lua: &Lua, mouse_table: &Table) -> rlua::Result<()> {
    mouse_table.set("__data", MouseState::default().to_lua(lua)?)
}

fn meta_setup(lua: &Lua, mouse_table: &Table) -> rlua::Result<()> {
    let meta_table = lua.create_table();
    meta_table.set("__tostring", lua.create_function(|_, val: Table| {
        Ok(format!("{}", val.get::<_, MouseState>("__data")?))
    }))?;
    mouse_table.set_metatable(Some(meta_table));
    Ok(())
}

fn method_setup(lua: &Lua, mouse_table: &Table) -> rlua::Result<()> {
    mouse_table.set("coords", lua.create_function(coords))?;
    Ok(())
}

fn coords<'lua>(lua: &'lua Lua, (coords, _ignore_enter): (rlua::Value<'lua>, rlua::Value<'lua>))
                -> rlua::Result<Table<'lua>> {
    match coords {
        rlua::Value::Table(coords) => {
            let (x, y) = (coords.get("x")?, coords.get("y")?);
            // TODO The ignore_enter is supposed to not send a send event to the client
            // That's not possible, at least until wlroots is complete.
            input::pointer::set_position(Point { x, y });
            Ok(coords)
        },
        _ => {
            // get the coords
            let coords = lua.create_table();
            let Point { x, y } = input::pointer::get_position();
            coords.set("x", x)?;
            coords.set("y", y)?;
            // TODO It expects a table of what buttons were pressed.
            coords.set("buttons", lua.create_table())?;
            Ok(coords)
        }
    }
}
<|endoftext|>