From c2702fa36b7a47268ac748b1fba911f9805419df Mon Sep 17 00:00:00 2001
From: shivshank <shivshank@users.noreply.github.com>
Date: Fri, 27 Jul 2018 20:15:56 -0400
Subject: [PATCH] Add safety check for update_buffer

---
 README.md                      |  2 --
 examples/wrong_sized_buffer.rs | 20 +++++++++++++++++++
 src/core.rs                    | 36 ++++++++++++++++++++++++++++++++--
 src/lib.rs                     |  5 +++++
 4 files changed, 59 insertions(+), 4 deletions(-)
 create mode 100644 examples/wrong_sized_buffer.rs

diff --git a/README.md b/README.md
index 88aa42f..1a41668 100644
--- a/README.md
+++ b/README.md
@@ -51,8 +51,6 @@ See the docs for more info.
 
 Listed in rough order of importance and ease (which are surprisingly correlated here!).
 
- - Bounds check on `update_buffer` which will currently segfault if you pass the wrong size.
-
  - Shader playground. Add a method for using shadertoy-like fragment shaders to be applied to
     your submitted pixel data.
 
diff --git a/examples/wrong_sized_buffer.rs b/examples/wrong_sized_buffer.rs
new file mode 100644
index 0000000..463154e
--- /dev/null
+++ b/examples/wrong_sized_buffer.rs
@@ -0,0 +1,20 @@
+extern crate mini_gl_fb;
+
+use mini_gl_fb::{Config, BufferFormat};
+
+fn main() {
+    let mut fb = mini_gl_fb::get_fancy(Config {
+        window_title: "Hello world!",
+        window_size: (800.0, 600.0),
+        buffer_size: (2, 2),
+        .. Default::default()
+    });
+
+    fb.change_buffer_format::<u8>(BufferFormat::RG);
+
+    // This should panic! We should only be providing two components but we provide 4!
+    let buffer = vec![[0u8, 50, 128, 255]; 4];
+    fb.update_buffer(&buffer);
+
+    fb.persist();
+}
diff --git a/src/core.rs b/src/core.rs
index 5b545ad..e4e1c9f 100644
--- a/src/core.rs
+++ b/src/core.rs
@@ -17,6 +17,7 @@ use gl;
 use gl::types::*;
 
 use std::ptr::null;
+use std::mem::size_of_val;
 
 /// Create a context using glutin given a configuration.
 pub fn init_glutin_context<S: ToString>(
@@ -216,10 +217,22 @@ pub struct Framebuffer {
 
 impl Framebuffer {
     pub fn update_buffer<T>(&mut self, image_data: &[T]) {
-        // TODO: Safety check on the length of the passed slice so this is actually a safe method
+        // Check the length of the passed slice so this is actually a safe method.
+        let (format, kind) = self.texture_format;
+        let expected_size_in_bytes = size_of_gl_type_enum(kind)
+            * format.components()
+            * self.buffer_width as usize
+            * self.buffer_height as usize;
+        let actual_size_in_bytes = size_of_val(image_data);
+        if actual_size_in_bytes != expected_size_in_bytes {
+            panic!(
+                "Expected a buffer of {} bytes, instead recieved one of {} bytes",
+                expected_size_in_bytes,
+                actual_size_in_bytes
+            );
+        }
         self.draw(|fb| {
             unsafe {
-                let (format, kind) = fb.texture_format;
                 gl::TexImage2D(
                     gl::TEXTURE_2D,
                     0,
@@ -317,6 +330,18 @@ pub enum BufferFormat {
     BGRA = gl::BGRA,
 }
 
+impl BufferFormat {
+    fn components(&self) -> usize {
+        use self::BufferFormat::*;
+        match self {
+            R => 1,
+            RG => 2,
+            RGB | BGR => 3,
+            RGBA | BGRA => 4,
+        }
+    }
+}
+
 pub trait ToGlType {
     fn to_gl_enum() -> GLenum;
 }
@@ -342,6 +367,13 @@ impl_ToGlType!(
     i8, gl::BYTE,
 );
 
+fn size_of_gl_type_enum(gl_enum: GLenum) -> usize {
+    match gl_enum {
+        gl::UNSIGNED_BYTE | gl::BYTE => 1,
+        _ => panic!("Must pass a GL enum representing a type"),
+    }
+}
+
 fn create_texture(width: i32, height: i32, format: BufferFormat, buffer_kind: GLenum) -> GLuint {
     unsafe {
         let mut tex = 0;
diff --git a/src/lib.rs b/src/lib.rs
index 15384dc..0477b30 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -160,6 +160,11 @@ impl MiniGlFb {
     /// Updates the backing buffer and draws immediately (swaps buffers).
     ///
     /// The main drawing function.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the size of the buffer does not exactly match the correct size of the texture
+    /// data required based on the buffers format.
     pub fn update_buffer<T>(&mut self, image_data: &[T]) {
         self.internal.update_buffer(image_data);
     }