9 files changed, 66 insertions, 43 deletions

diff --git a/minui/graphics_adf.h b/minui/graphics_adf.h
index bf9842878..79d8d2acb 100644
--- a/minui/graphics_adf.h
+++ b/minui/graphics_adf.h
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <stddef.h>
 #include <stdint.h>
 #include <sys/types.h>
 
@@ -40,8 +41,8 @@ class GRSurfaceAdf : public GRSurface {
  private:
   friend class MinuiBackendAdf;
 
-  GRSurfaceAdf(int width, int height, int row_bytes, int pixel_bytes, __u32 offset, __u32 pitch,
-               int fd)
+  GRSurfaceAdf(size_t width, size_t height, size_t row_bytes, size_t pixel_bytes, __u32 offset,
+               __u32 pitch, int fd)
       : GRSurface(width, height, row_bytes, pixel_bytes), offset(offset), pitch(pitch), fd(fd) {}
 
   const __u32 offset;
diff --git a/minui/graphics_drm.h b/minui/graphics_drm.h
index 6ba46e60b..57ba39b83 100644
--- a/minui/graphics_drm.h
+++ b/minui/graphics_drm.h
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <stddef.h>
 #include <stdint.h>
 
 #include <memory>
@@ -39,7 +40,8 @@ class GRSurfaceDrm : public GRSurface {
  private:
   friend class MinuiBackendDrm;
 
-  GRSurfaceDrm(int width, int height, int row_bytes, int pixel_bytes, int drm_fd, uint32_t handle)
+  GRSurfaceDrm(size_t width, size_t height, size_t row_bytes, size_t pixel_bytes, int drm_fd,
+               uint32_t handle)
       : GRSurface(width, height, row_bytes, pixel_bytes), drm_fd_(drm_fd), handle(handle) {}
 
   const int drm_fd_;
diff --git a/minui/graphics_fbdev.cpp b/minui/graphics_fbdev.cpp
index 93e4420d3..8d9c9741d 100644
--- a/minui/graphics_fbdev.cpp
+++ b/minui/graphics_fbdev.cpp
@@ -32,8 +32,8 @@
 
 #include "minui/minui.h"
 
-std::unique_ptr<GRSurfaceFbdev> GRSurfaceFbdev::Create(int width, int height, int row_bytes,
-                                                       int pixel_bytes) {
+std::unique_ptr<GRSurfaceFbdev> GRSurfaceFbdev::Create(size_t width, size_t height,
+                                                       size_t row_bytes, size_t pixel_bytes) {
   // Cannot use std::make_unique to access non-public ctor.
   return std::unique_ptr<GRSurfaceFbdev>(new GRSurfaceFbdev(width, height, row_bytes, pixel_bytes));
 }
@@ -130,7 +130,7 @@ GRSurface* MinuiBackendFbdev::Init() {
   fb_fd = std::move(fd);
   SetDisplayedFramebuffer(0);
 
-  printf("framebuffer: %d (%d x %d)\n", fb_fd.get(), gr_draw->width, gr_draw->height);
+  printf("framebuffer: %d (%zu x %zu)\n", fb_fd.get(), gr_draw->width, gr_draw->height);
 
   Blank(true);
   Blank(false);
diff --git a/minui/graphics_fbdev.h b/minui/graphics_fbdev.h
index 016ab88bc..596ba74ea 100644
--- a/minui/graphics_fbdev.h
+++ b/minui/graphics_fbdev.h
@@ -17,6 +17,7 @@
 #pragma once
 
 #include <linux/fb.h>
+#include <stddef.h>
 #include <stdint.h>
 
 #include <memory>
@@ -30,8 +31,8 @@
 class GRSurfaceFbdev : public GRSurface {
  public:
   // Creates and returns a GRSurfaceFbdev instance, or nullptr on error.
-  static std::unique_ptr<GRSurfaceFbdev> Create(int width, int height, int row_bytes,
-                                                int pixel_bytes);
+  static std::unique_ptr<GRSurfaceFbdev> Create(size_t width, size_t height, size_t row_bytes,
+                                                size_t pixel_bytes);
 
   uint8_t* data() override {
     return buffer_;
diff --git a/minui/include/minui/minui.h b/minui/include/minui/minui.h
index 3231248a0..e49c6ac97 100644
--- a/minui/include/minui/minui.h
+++ b/minui/include/minui/minui.h
@@ -33,13 +33,16 @@
 
 class GRSurface {
  public:
+  static constexpr size_t kSurfaceDataAlignment = 8;
+
   virtual ~GRSurface() = default;
 
   // Creates and returns a GRSurface instance that's sufficient for storing an image of the given
-  // size. The starting address of the surface data is aligned to SURFACE_DATA_ALIGNMENT. Returns
-  // the created GRSurface instance (in std::unique_ptr), or nullptr on error.
-  static std::unique_ptr<GRSurface> Create(int width, int height, int row_bytes, int pixel_bytes,
-                                           size_t data_size);
+  // size (i.e. row_bytes * height). The starting address of the surface data is aligned to
+  // kSurfaceDataAlignment. Returns the created GRSurface instance (in std::unique_ptr), or nullptr
+  // on error.
+  static std::unique_ptr<GRSurface> Create(size_t width, size_t height, size_t row_bytes,
+                                           size_t pixel_bytes);
 
   // Clones the current GRSurface instance (i.e. an image).
   std::unique_ptr<GRSurface> Clone() const;
@@ -52,13 +55,17 @@ class GRSurface {
     return const_cast<const uint8_t*>(const_cast<GRSurface*>(this)->data());
   }
 
-  int width;
-  int height;
-  int row_bytes;
-  int pixel_bytes;
+  size_t data_size() const {
+    return data_size_;
+  }
+
+  size_t width;
+  size_t height;
+  size_t row_bytes;
+  size_t pixel_bytes;
 
  protected:
-  GRSurface(int width, int height, int row_bytes, int pixel_bytes)
+  GRSurface(size_t width, size_t height, size_t row_bytes, size_t pixel_bytes)
       : width(width), height(height), row_bytes(row_bytes), pixel_bytes(pixel_bytes) {}
 
  private:
diff --git a/minui/resources.cpp b/minui/resources.cpp
index 057d3fd03..069a49529 100644
--- a/minui/resources.cpp
+++ b/minui/resources.cpp
@@ -27,6 +27,7 @@
 #include <sys/types.h>
 #include <unistd.h>
 
+#include <limits>
 #include <memory>
 #include <regex>
 #include <string>
@@ -39,11 +40,14 @@
 
 static std::string g_resource_dir{ "/res/images" };
 
-std::unique_ptr<GRSurface> GRSurface::Create(int width, int height, int row_bytes, int pixel_bytes,
-                                             size_t data_size) {
-  static constexpr size_t kSurfaceDataAlignment = 8;
+std::unique_ptr<GRSurface> GRSurface::Create(size_t width, size_t height, size_t row_bytes,
+                                             size_t pixel_bytes) {
+  if (width == 0 || row_bytes == 0 || height == 0 || pixel_bytes == 0) return nullptr;
+  if (std::numeric_limits<size_t>::max() / row_bytes < height) return nullptr;
+
   // Cannot use std::make_unique to access non-public ctor.
   auto result = std::unique_ptr<GRSurface>(new GRSurface(width, height, row_bytes, pixel_bytes));
+  size_t data_size = row_bytes * height;
   result->data_size_ =
       (data_size + kSurfaceDataAlignment - 1) / kSurfaceDataAlignment * kSurfaceDataAlignment;
   result->data_.reset(
@@ -53,7 +57,7 @@ std::unique_ptr<GRSurface> GRSurface::Create(int width, int height, int row_byte
 }
 
 std::unique_ptr<GRSurface> GRSurface::Clone() const {
-  auto result = GRSurface::Create(width, height, row_bytes, pixel_bytes, data_size_);
+  auto result = GRSurface::Create(width, height, row_bytes, pixel_bytes);
   if (!result) return nullptr;
   memcpy(result->data(), data(), data_size_);
   return result;
@@ -189,7 +193,7 @@ int res_create_display_surface(const char* name, GRSurface** pSurface) {
   png_uint_32 width = png_handler.width();
   png_uint_32 height = png_handler.height();
 
-  auto surface = GRSurface::Create(width, height, width * 4, 4, width * height * 4);
+  auto surface = GRSurface::Create(width, height, width * 4, 4);
   if (!surface) {
     return -8;
   }
@@ -259,9 +263,7 @@ int res_create_multi_display_surface(const char* name, int* frames, int* fps,
     goto exit;
   }
   for (int i = 0; i < *frames; ++i) {
-    auto height_per_frame = height / *frames;
-    auto created_surface =
-        GRSurface::Create(width, height_per_frame, width * 4, 4, width * height_per_frame * 4);
+    auto created_surface = GRSurface::Create(width, height / *frames, width * 4, 4);
     if (!created_surface) {
       result = -8;
       goto exit;
@@ -309,7 +311,7 @@ int res_create_alpha_surface(const char* name, GRSurface** pSurface) {
   png_uint_32 width = png_handler.width();
   png_uint_32 height = png_handler.height();
 
-  auto surface = GRSurface::Create(width, height, width, 1, width * height);
+  auto surface = GRSurface::Create(width, height, width, 1);
   if (!surface) {
     return -8;
   }
@@ -415,7 +417,7 @@ int res_create_localized_alpha_surface(const char* name,
     if (y + 1 + h >= height || matches_locale(loc, locale)) {
       printf("  %20s: %s (%d x %d @ %d)\n", name, loc, w, h, y);
 
-      auto surface = GRSurface::Create(w, h, w, 1, w * h);
+      auto surface = GRSurface::Create(w, h, w, 1);
       if (!surface) {
         return -8;
       }
diff --git a/screen_ui.cpp b/screen_ui.cpp
index ed71888d1..765d2fe60 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -282,14 +282,14 @@ bool GraphicMenu::ValidateGraphicSurface(size_t max_width, size_t max_height, in
   }
 
   if (surface->pixel_bytes != 1 || surface->width != surface->row_bytes) {
-    fprintf(stderr, "Invalid graphic surface, pixel bytes: %d, width: %d row_bytes: %d",
+    fprintf(stderr, "Invalid graphic surface, pixel bytes: %zu, width: %zu row_bytes: %zu",
             surface->pixel_bytes, surface->width, surface->row_bytes);
     return false;
   }
 
   if (surface->width > max_width || surface->height > max_height - y) {
     fprintf(stderr,
-            "Graphic surface doesn't fit into the screen. width: %d, height: %d, max_width: %zu,"
+            "Graphic surface doesn't fit into the screen. width: %zu, height: %zu, max_width: %zu,"
             " max_height: %zu, vertical offset: %d\n",
             surface->width, surface->height, max_width, max_height, y);
     return false;
diff --git a/tests/unit/minui_test.cpp b/tests/unit/minui_test.cpp
index d68e5e3a1..c7d7f7eef 100644
--- a/tests/unit/minui_test.cpp
+++ b/tests/unit/minui_test.cpp
@@ -17,6 +17,7 @@
 #include <stdint.h>
 #include <stdlib.h>
 
+#include <limits>
 #include <vector>
 
 #include <gtest/gtest.h>
@@ -24,21 +25,30 @@
 #include "minui/minui.h"
 
 TEST(GRSurfaceTest, Create_aligned) {
-  static constexpr size_t kSurfaceDataAlignment = 8;
-  for (size_t data_size = 100; data_size < 128; data_size++) {
-    auto surface = GRSurface::Create(10, 1, 10, 1, data_size);
-    ASSERT_TRUE(surface);
-    ASSERT_EQ(0, reinterpret_cast<uintptr_t>(surface->data()) % kSurfaceDataAlignment);
-  }
+  auto surface = GRSurface::Create(9, 11, 9, 1);
+  ASSERT_TRUE(surface);
+  ASSERT_EQ(0, reinterpret_cast<uintptr_t>(surface->data()) % GRSurface::kSurfaceDataAlignment);
+  // data_size will be rounded up to the next multiple of GRSurface::kSurfaceDataAlignment.
+  ASSERT_EQ(0, surface->data_size() % GRSurface::kSurfaceDataAlignment);
+  ASSERT_GE(surface->data_size(), 11 * 9);
+}
+
+TEST(GRSurfaceTest, Create_invalid_inputs) {
+  ASSERT_FALSE(GRSurface::Create(9, 11, 0, 1));
+  ASSERT_FALSE(GRSurface::Create(9, 0, 9, 1));
+  ASSERT_FALSE(GRSurface::Create(0, 11, 9, 1));
+  ASSERT_FALSE(GRSurface::Create(9, 11, 9, 0));
+  ASSERT_FALSE(GRSurface::Create(9, 101, std::numeric_limits<size_t>::max() / 100, 1));
 }
 
 TEST(GRSurfaceTest, Clone) {
-  static constexpr size_t kImageSize = 10 * 50;
-  auto image = GRSurface::Create(50, 10, 50, 1, kImageSize);
-  for (auto i = 0; i < kImageSize; i++) {
+  auto image = GRSurface::Create(50, 10, 50, 1);
+  ASSERT_GE(image->data_size(), 10 * 50);
+  for (auto i = 0; i < image->data_size(); i++) {
     image->data()[i] = rand() % 128;
   }
   auto image_copy = image->Clone();
-  ASSERT_EQ(std::vector(image->data(), image->data() + kImageSize),
-            std::vector(image_copy->data(), image_copy->data() + kImageSize));
+  ASSERT_EQ(image->data_size(), image_copy->data_size());
+  ASSERT_EQ(std::vector(image->data(), image->data() + image->data_size()),
+            std::vector(image_copy->data(), image_copy->data() + image->data_size()));
 }
diff --git a/tests/unit/screen_ui_test.cpp b/tests/unit/screen_ui_test.cpp
index 09c49977f..61a092551 100644
--- a/tests/unit/screen_ui_test.cpp
+++ b/tests/unit/screen_ui_test.cpp
@@ -231,7 +231,7 @@ TEST_F(ScreenUITest, WearMenuSelectItemsOverflow) {
 }
 
 TEST_F(ScreenUITest, GraphicMenuSelection) {
-  auto image = GRSurface::Create(50, 50, 50, 1, 50 * 50);
+  auto image = GRSurface::Create(50, 50, 50, 1);
   auto header = image->Clone();
   std::vector<const GRSurface*> items = {
     image.get(),
@@ -258,7 +258,7 @@ TEST_F(ScreenUITest, GraphicMenuSelection) {
 }
 
 TEST_F(ScreenUITest, GraphicMenuValidate) {
-  auto image = GRSurface::Create(50, 50, 50, 1, 50 * 50);
+  auto image = GRSurface::Create(50, 50, 50, 1);
   auto header = image->Clone();
   std::vector<const GRSurface*> items = {
     image.get(),
@@ -269,7 +269,7 @@ TEST_F(ScreenUITest, GraphicMenuValidate) {
   ASSERT_TRUE(GraphicMenu::Validate(200, 200, header.get(), items));
 
   // Menu exceeds the horizontal boundary.
-  auto wide_surface = GRSurface::Create(300, 50, 300, 1, 300 * 50);
+  auto wide_surface = GRSurface::Create(300, 50, 300, 1);
   ASSERT_FALSE(GraphicMenu::Validate(299, 200, wide_surface.get(), items));
 
   // Menu exceeds the vertical boundary.