Remove panic (#6)

* remove panic and return image length 0 in header on error

* Fix index error, add pry

* Run resize in a thread so ruby will GC it

Ruby does not GC the wasm instance in creators unless we run this in a
thread and extract the value.

Also added a test for resizing a photo multiple times in order to test
crashing, memory use, etc.

* Disable GC while calling WASM

Puma crashes unless we disable GC while interacting with wasmer.

* Update python lib to not use exceptions

* Update thumbnail.py

Co-authored-by: eV <[email protected]>
Co-authored-by: eV <[email protected]>

wasm-thumbnail-py/src/wasm_thumbnail/thumbnail.py

@@ -7,13 +7,17 @@
 from wasmer import engine, Store, Module, Instance, ImportObject, Function, FunctionType, Type
 from wasmer_compiler_cranelift import Compiler
 
-def decode_padded_image(data):
-    """Extract a payload from its padding by reading its length header."""
+def get_unpadded_length(data):
+    """Get the unpadded length from the header."""
     data_length_without_header = len(data) - 4
     if data_length_without_header < 0:
         raise ValueError('Data must be at least 4 bytes long', len(data))
 
-    payload_length = struct.unpack('!L', data[0:4])[0]
+    return struct.unpack('!L', data[0:4])[0]
+
+def decode_padded_image(data):
+    """Extract a payload from its padding by reading its length header."""
+    payload_length = get_unpadded_length(data)
 
     if data_length_without_header < payload_length:
         raise ValueError('Payload is shorter than the expected length',
@@ -38,6 +42,10 @@ def resize_and_pad_image(image_bytes, width, height, size, quality = 80):
     out_bytes = bytes(memory[:size])
     instance.exports.deallocate(output_pointer, size)
 
+    unpadded_length = get_unpadded_length(out_bytes)
+    if unpadded_length == 0:
+        raise RuntimeError('Image resizing failed')
+
     return out_bytes
 
 def register_panic(msg_ptr: int, msg_len: int, file_ptr: int, file_len: int, line: int, column: int):

wasm-thumbnail-rb/lib/wasm/thumbnail/rb.rb

@@ -14,55 +14,69 @@ def self.register_panic(_msg_ptr = nil, _msg_len = nil, _file_ptr = nil, _file_l
       end
 
       def self.resize_and_pad_with_header(file_bytes:, width:, height:, size:, quality: 80)
-        # Let's compile the module to be able to execute it!
-        wasm_instance = Wasm::Thumbnail::Rb::GetWasmInstance.call
+        # Ruby does not garbage collect the wasm instance unless we run this in a thread
+        t = Thread.new do
+          Thread.current.abort_on_exception = true
 
-        # This tells us how much space we'll need to put our image in the WASM env
-        image_length = file_bytes.length
-        input_pointer = wasm_instance.exports.allocate.call(image_length)
-        # Get a pointer on the allocated memory so we can write to it
-        memory = wasm_instance.exports.memory.uint8_view input_pointer
+          # Let's compile the module to be able to execute it!
+          wasm_instance = Wasm::Thumbnail::Rb::GetWasmInstance.call
 
-        # Put the image to resize in the allocated space
-        (0..image_length - 1).each do |nth|
-          memory[nth] = file_bytes[nth]
-        end
+          # This tells us how much space we'll need to put our image in the WASM env
+          image_length = file_bytes.length
+          input_pointer = wasm_instance.exports.allocate.call(image_length)
+          # Get a pointer on the allocated memory so we can write to it
+          memory = wasm_instance.exports.memory.uint8_view input_pointer
+
+          # Put the image to resize in the allocated space
+          (0..image_length - 1).each do |nth|
+            memory[nth] = file_bytes[nth]
+          end
+
+          # Do the actual resize and pad
+          # Note that this writes to a portion of memory the new JPEG file, but right pads the rest of the space
+          # we gave it with 0.
+          begin
+            output_pointer = wasm_instance.exports.resize_and_pad.call(input_pointer,
+                                                                      image_length,
+                                                                      width,
+                                                                      height,
+                                                                      size,
+                                                                      quality)
+          rescue RuntimeError => e
+            # Deallocate
+            wasm_instance.exports.deallocate.call(input_pointer, image_length)
+            raise "Error processing the image: #{e.message}"
+          end
+          # Get a pointer to the result
+          memory = wasm_instance.exports.memory.uint8_view output_pointer
+
+          # Only take the buffer that we told the rust function we needed. The resize function
+          # makes a smaller image than the buffer we said, and then pads out the rest.
+          bytes = memory.to_a.take(size)
 
-        # Do the actual resize and pad
-        # Note that this writes to a portion of memory the new JPEG file, but right pads the rest of the space
-        # we gave it with 0.
-        begin
-          output_pointer = wasm_instance.exports.resize_and_pad.call(input_pointer,
-                                                                     image_length,
-                                                                     width,
-                                                                     height,
-                                                                     size,
-                                                                     quality)
-        rescue RuntimeError
           # Deallocate
           wasm_instance.exports.deallocate.call(input_pointer, image_length)
-          raise "Error processing the image."
-        end
-        # Get a pointer to the result
-        memory = wasm_instance.exports.memory.uint8_view output_pointer
-
-        # Only take the buffer that we told the rust function we needed. The resize function
-        # makes a smaller image than the buffer we said, and then pads out the rest.
-        bytes = memory.to_a.take(size)
+          wasm_instance.exports.deallocate.call(output_pointer, size)
 
-        # Deallocate
-        wasm_instance.exports.deallocate.call(input_pointer, image_length)
-        wasm_instance.exports.deallocate.call(output_pointer, size)
+          bytes
+        end
 
-        bytes
+        t.join
+        t.value
       end
 
       def self.resize_and_pad(file_bytes:, width:, height:, size:, quality: 80)
+        GC.disable
         bytes = resize_and_pad_with_header(file_bytes: file_bytes,
                                            width: width,
                                            height: height,
                                            size: size + 4,
                                            quality: quality)
+        GC.enable
+        len = bytes[..3].pack("cccc").unpack("N*")[0]
+        if len == 0
+          raise "Error processing the image."
+        end
 
         # The first 4 bytes are a header until the image. The actual image probably ends well before
         # the whole buffer, but we keep the junk data on the end to make all the images the same size

wasm-thumbnail-rb/test/wasm/thumbnail/rb_test.rb

@@ -36,4 +36,26 @@ def test_it_calls_panic_if_image_too_big
     end
     assert_equal(exception.message, "Error processing the image.")
   end
+
+  def test_should_not_crash_after_many_resizes
+    # Photo from https://unsplash.com/photos/SRFhLQ_zmak
+    file_bytes = File.binread("#{__dir__}/shipwreck.jpg").unpack("C*")
+    quality_to_try = 100
+    # Run a few times. Adjust as needed, checking memory usage as well
+    while quality_to_try > 96
+      begin
+        return Wasm::Thumbnail::Rb.resize_and_pad(
+          file_bytes: file_bytes,
+          width: 2700,
+          height: 528,
+          size: 5,
+          quality: quality_to_try
+        )
+      rescue RuntimeError => e
+        quality_to_try -= 1
+        puts "Trying quality #{quality_to_try}"
+        next
+      end
+    end
+  end
 end

wasm-thumbnail-rb/wasm-thumbnail-rb.gemspec

@@ -25,4 +25,7 @@ Gem::Specification.new do |spec|
   spec.require_paths = ["lib"]
 
   spec.add_dependency "wasmer", "~> 1.0"
+
+  spec.add_development_dependency "pry"
+
 end

wasm-thumbnail/src/lib.rs

@@ -1,3 +1,4 @@
+use std::error::Error;
 use std::mem;
 use std::os::raw::c_void;
 
@@ -21,36 +22,19 @@ pub extern "C" fn resize_and_pad(
     nwidth: u32,
     nheight: u32,
     nsize: usize,
-    nquality: u8
+    nquality: u8,
 ) -> *const u8 {
     register_panic_hook();
 
     let slice: &[u8] = unsafe { std::slice::from_raw_parts(pointer, length) };
 
-    let img = image::load_from_memory(slice).expect("must be a valid image");
-
-    // Resize preserves aspect ratio
-    let img = img.resize(nwidth, nheight, FilterType::Lanczos3);
-
-    // Copy pixels only
-    let mut result = DynamicImage::new_rgba8(img.width(), img.height());
-    result
-        .copy_from(&img, 0, 0)
-        .expect("copy should not fail as output is same dimensions");
-
     let mut out: Vec<u8> = Vec::with_capacity(nsize);
-
     // Reserve space at the start for length header
     out.extend_from_slice(&[0, 0, 0, 0]);
 
-    result
-        .write_to(&mut out, ImageOutputFormat::Jpeg(nquality))
-        .expect("can save as jpeg");
-
-    assert!(out.len() <= nsize);
-
-    let thumbnail_len: u32 = out.len() as u32 - 4;
-    out.splice(..4, thumbnail_len.to_be_bytes().iter().cloned());
+    if let Ok(thumbnail_len) = _resize_and_pad(slice, &mut out, nwidth, nheight, nsize, nquality) {
+        out.splice(..4, thumbnail_len.to_be_bytes().iter().cloned());
+    }
 
     out.resize(nsize, 0);
 
@@ -59,6 +43,32 @@ pub extern "C" fn resize_and_pad(
     pointer
 }
 
+pub fn _resize_and_pad(
+    slice: &[u8],
+    out: &mut Vec<u8>,
+    nwidth: u32,
+    nheight: u32,
+    nsize: usize,
+    nquality: u8,
+) -> Result<u32, Box<dyn Error>> {
+    let img = image::load_from_memory(slice)?;
+
+    // Resize preserves aspect ratio
+    let img = img.resize(nwidth, nheight, FilterType::Lanczos3);
+
+    // Copy pixels only
+    let mut result = DynamicImage::new_rgba8(img.width(), img.height());
+    result.copy_from(&img, 0, 0)?;
+
+    result.write_to(out, ImageOutputFormat::Jpeg(nquality))?;
+
+    if out.len() > nsize {
+        return Err("size is too large".into());
+    }
+
+    Ok(out.len() as u32 - 4)
+}
+
 /// Allocate a new buffer in the wasm memory space
 #[no_mangle]
 pub extern "C" fn allocate(capacity: usize) -> *mut c_void {