19 files changed, 402 insertions, 186 deletions

diff --git a/install.cpp b/install.cpp
index e379ef307..42d264157 100644
--- a/install.cpp
+++ b/install.cpp
@@ -695,18 +695,18 @@ int install_package(const std::string& path, bool* wipe_cache, bool needs_mount,
 }
 
 bool verify_package(const unsigned char* package_data, size_t package_size) {
-  static constexpr const char* PUBLIC_KEYS_FILE = "/res/keys";
-  std::vector<Certificate> loadedKeys;
-  if (!load_keys(PUBLIC_KEYS_FILE, loadedKeys)) {
+  static constexpr const char* CERTIFICATE_ZIP_FILE = "/system/etc/security/otacerts.zip";
+  std::vector<Certificate> loaded_keys = LoadKeysFromZipfile(CERTIFICATE_ZIP_FILE);
+  if (loaded_keys.empty()) {
     LOG(ERROR) << "Failed to load keys";
     return false;
   }
-  LOG(INFO) << loadedKeys.size() << " key(s) loaded from " << PUBLIC_KEYS_FILE;
+  LOG(INFO) << loaded_keys.size() << " key(s) loaded from " << CERTIFICATE_ZIP_FILE;
 
   // Verify package.
   ui->Print("Verifying update package...\n");
   auto t0 = std::chrono::system_clock::now();
-  int err = verify_file(package_data, package_size, loadedKeys,
+  int err = verify_file(package_data, package_size, loaded_keys,
                         std::bind(&RecoveryUI::SetProgress, ui, std::placeholders::_1));
   std::chrono::duration<double> duration = std::chrono::system_clock::now() - t0;
   ui->Print("Update package verification took %.1f s (result %d).\n", duration.count(), err);
diff --git a/minui/graphics.cpp b/minui/graphics.cpp
index e6367d950..4d1f9b2d2 100644
--- a/minui/graphics.cpp
+++ b/minui/graphics.cpp
@@ -40,7 +40,7 @@ static uint32_t gr_current = ~0;
 static constexpr uint32_t alpha_mask = 0xff000000;
 
 // gr_draw is owned by backends.
-static const GRSurface* gr_draw = nullptr;
+static GRSurface* gr_draw = nullptr;
 static GRRotation rotation = GRRotation::NONE;
 static PixelFormat pixel_format = PixelFormat::UNKNOWN;
 
@@ -121,28 +121,29 @@ static void incr_y(uint32_t** p, int row_pixels) {
 }
 
 // Returns pixel pointer at given coordinates with rotation adjustment.
-static uint32_t* pixel_at(const GRSurface* surf, int x, int y, int row_pixels) {
+static uint32_t* PixelAt(GRSurface* surface, int x, int y, int row_pixels) {
   switch (rotation) {
     case GRRotation::NONE:
-      return reinterpret_cast<uint32_t*>(surf->data) + y * row_pixels + x;
+      return reinterpret_cast<uint32_t*>(surface->data()) + y * row_pixels + x;
     case GRRotation::RIGHT:
-      return reinterpret_cast<uint32_t*>(surf->data) + x * row_pixels + (surf->width - y);
+      return reinterpret_cast<uint32_t*>(surface->data()) + x * row_pixels + (surface->width - y);
     case GRRotation::DOWN:
-      return reinterpret_cast<uint32_t*>(surf->data) + (surf->height - 1 - y) * row_pixels +
-             (surf->width - 1 - x);
+      return reinterpret_cast<uint32_t*>(surface->data()) + (surface->height - 1 - y) * row_pixels +
+             (surface->width - 1 - x);
     case GRRotation::LEFT:
-      return reinterpret_cast<uint32_t*>(surf->data) + (surf->height - 1 - x) * row_pixels + y;
+      return reinterpret_cast<uint32_t*>(surface->data()) + (surface->height - 1 - x) * row_pixels +
+             y;
     default:
       printf("invalid rotation %d", static_cast<int>(rotation));
   }
   return nullptr;
 }
 
-static void text_blend(uint8_t* src_p, int src_row_bytes, uint32_t* dst_p, int dst_row_pixels,
-                       int width, int height) {
+static void TextBlend(const uint8_t* src_p, int src_row_bytes, uint32_t* dst_p, int dst_row_pixels,
+                      int width, int height) {
   uint8_t alpha_current = static_cast<uint8_t>((alpha_mask & gr_current) >> 24);
   for (int j = 0; j < height; ++j) {
-    uint8_t* sx = src_p;
+    const uint8_t* sx = src_p;
     uint32_t* px = dst_p;
     for (int i = 0; i < width; ++i, incr_x(&px, dst_row_pixels)) {
       uint8_t a = *sx++;
@@ -176,18 +177,18 @@ void gr_text(const GRFont* font, int x, int y, const char* s, bool bold) {
     }
 
     int row_pixels = gr_draw->row_bytes / gr_draw->pixel_bytes;
-    uint8_t* src_p = font->texture->data + ((ch - ' ') * font->char_width) +
-                     (bold ? font->char_height * font->texture->row_bytes : 0);
-    uint32_t* dst_p = pixel_at(gr_draw, x, y, row_pixels);
+    const uint8_t* src_p = font->texture->data() + ((ch - ' ') * font->char_width) +
+                           (bold ? font->char_height * font->texture->row_bytes : 0);
+    uint32_t* dst_p = PixelAt(gr_draw, x, y, row_pixels);
 
-    text_blend(src_p, font->texture->row_bytes, dst_p, row_pixels, font->char_width,
-               font->char_height);
+    TextBlend(src_p, font->texture->row_bytes, dst_p, row_pixels, font->char_width,
+              font->char_height);
 
     x += font->char_width;
   }
 }
 
-void gr_texticon(int x, int y, GRSurface* icon) {
+void gr_texticon(int x, int y, const GRSurface* icon) {
   if (icon == nullptr) return;
 
   if (icon->pixel_bytes != 1) {
@@ -201,10 +202,9 @@ void gr_texticon(int x, int y, GRSurface* icon) {
   if (outside(x, y) || outside(x + icon->width - 1, y + icon->height - 1)) return;
 
   int row_pixels = gr_draw->row_bytes / gr_draw->pixel_bytes;
-  uint8_t* src_p = icon->data;
-  uint32_t* dst_p = pixel_at(gr_draw, x, y, row_pixels);
-
-  text_blend(src_p, icon->row_bytes, dst_p, row_pixels, icon->width, icon->height);
+  const uint8_t* src_p = icon->data();
+  uint32_t* dst_p = PixelAt(gr_draw, x, y, row_pixels);
+  TextBlend(src_p, icon->row_bytes, dst_p, row_pixels, icon->width, icon->height);
 }
 
 void gr_color(unsigned char r, unsigned char g, unsigned char b, unsigned char a) {
@@ -221,9 +221,9 @@ void gr_clear() {
       (gr_current & 0xff) == ((gr_current >> 16) & 0xff) &&
       (gr_current & 0xff) == ((gr_current >> 24) & 0xff) &&
       gr_draw->row_bytes == gr_draw->width * gr_draw->pixel_bytes) {
-    memset(gr_draw->data, gr_current & 0xff, gr_draw->height * gr_draw->row_bytes);
+    memset(gr_draw->data(), gr_current & 0xff, gr_draw->height * gr_draw->row_bytes);
   } else {
-    uint32_t* px = reinterpret_cast<uint32_t*>(gr_draw->data);
+    uint32_t* px = reinterpret_cast<uint32_t*>(gr_draw->data());
     int row_diff = gr_draw->row_bytes / gr_draw->pixel_bytes - gr_draw->width;
     for (int y = 0; y < gr_draw->height; ++y) {
       for (int x = 0; x < gr_draw->width; ++x) {
@@ -244,7 +244,7 @@ void gr_fill(int x1, int y1, int x2, int y2) {
   if (outside(x1, y1) || outside(x2 - 1, y2 - 1)) return;
 
   int row_pixels = gr_draw->row_bytes / gr_draw->pixel_bytes;
-  uint32_t* p = pixel_at(gr_draw, x1, y1, row_pixels);
+  uint32_t* p = PixelAt(gr_draw, x1, y1, row_pixels);
   uint8_t alpha = static_cast<uint8_t>(((gr_current & alpha_mask) >> 24));
   if (alpha > 0) {
     for (int y = y1; y < y2; ++y) {
@@ -258,7 +258,7 @@ void gr_fill(int x1, int y1, int x2, int y2) {
   }
 }
 
-void gr_blit(GRSurface* source, int sx, int sy, int w, int h, int dx, int dy) {
+void gr_blit(const GRSurface* source, int sx, int sy, int w, int h, int dx, int dy) {
   if (source == nullptr) return;
 
   if (gr_draw->pixel_bytes != source->pixel_bytes) {
@@ -274,11 +274,12 @@ void gr_blit(GRSurface* source, int sx, int sy, int w, int h, int dx, int dy) {
   if (rotation != GRRotation::NONE) {
     int src_row_pixels = source->row_bytes / source->pixel_bytes;
     int row_pixels = gr_draw->row_bytes / gr_draw->pixel_bytes;
-    uint32_t* src_py = reinterpret_cast<uint32_t*>(source->data) + sy * source->row_bytes / 4 + sx;
-    uint32_t* dst_py = pixel_at(gr_draw, dx, dy, row_pixels);
+    const uint32_t* src_py =
+        reinterpret_cast<const uint32_t*>(source->data()) + sy * source->row_bytes / 4 + sx;
+    uint32_t* dst_py = PixelAt(gr_draw, dx, dy, row_pixels);
 
     for (int y = 0; y < h; y += 1) {
-      uint32_t* src_px = src_py;
+      const uint32_t* src_px = src_py;
       uint32_t* dst_px = dst_py;
       for (int x = 0; x < w; x += 1) {
         *dst_px = *src_px++;
@@ -288,8 +289,8 @@ void gr_blit(GRSurface* source, int sx, int sy, int w, int h, int dx, int dy) {
       incr_y(&dst_py, row_pixels);
     }
   } else {
-    unsigned char* src_p = source->data + sy * source->row_bytes + sx * source->pixel_bytes;
-    unsigned char* dst_p = gr_draw->data + dy * gr_draw->row_bytes + dx * gr_draw->pixel_bytes;
+    const uint8_t* src_p = source->data() + sy * source->row_bytes + sx * source->pixel_bytes;
+    uint8_t* dst_p = gr_draw->data() + dy * gr_draw->row_bytes + dx * gr_draw->pixel_bytes;
 
     for (int i = 0; i < h; ++i) {
       memcpy(dst_p, src_p, w * source->pixel_bytes);
diff --git a/minui/graphics_adf.cpp b/minui/graphics_adf.cpp
index 7439df9ac..6fc193f74 100644
--- a/minui/graphics_adf.cpp
+++ b/minui/graphics_adf.cpp
@@ -45,14 +45,14 @@ int MinuiBackendAdf::SurfaceInit(const drm_mode_modeinfo* mode, GRSurfaceAdf* su
   surf->row_bytes = surf->pitch;
   surf->pixel_bytes = (format == DRM_FORMAT_RGB565) ? 2 : 4;
 
-  surf->data = static_cast<uint8_t*>(
-      mmap(nullptr, surf->pitch * surf->height, PROT_WRITE, MAP_SHARED, surf->fd, surf->offset));
-  if (surf->data == MAP_FAILED) {
+  auto mmapped =
+      mmap(nullptr, surf->pitch * surf->height, PROT_WRITE, MAP_SHARED, surf->fd, surf->offset);
+  if (mmapped == MAP_FAILED) {
     int saved_errno = errno;
     close(surf->fd);
     return -saved_errno;
   }
-
+  surf->mmapped_buffer_ = static_cast<uint8_t*>(mmapped);
   return 0;
 }
 
@@ -185,7 +185,9 @@ void MinuiBackendAdf::Blank(bool blank) {
 }
 
 void MinuiBackendAdf::SurfaceDestroy(GRSurfaceAdf* surf) {
-  munmap(surf->data, surf->pitch * surf->height);
+  if (surf->mmapped_buffer_) {
+    munmap(surf->mmapped_buffer_, surf->pitch * surf->height);
+  }
   close(surf->fence_fd);
   close(surf->fd);
 }
diff --git a/minui/graphics_adf.h b/minui/graphics_adf.h
index 2f019ed0b..099d32962 100644
--- a/minui/graphics_adf.h
+++ b/minui/graphics_adf.h
@@ -14,21 +14,30 @@
  * limitations under the License.
  */
 
-#ifndef _GRAPHICS_ADF_H_
-#define _GRAPHICS_ADF_H_
+#pragma once
+
+#include <stdint.h>
 
 #include <adf/adf.h>
 
 #include "graphics.h"
+#include "minui/minui.h"
 
 class GRSurfaceAdf : public GRSurface {
+ public:
+  uint8_t* data() override {
+    return mmapped_buffer_;
+  }
+
  private:
+  friend class MinuiBackendAdf;
+
   int fence_fd;
   int fd;
   __u32 offset;
   __u32 pitch;
 
-  friend class MinuiBackendAdf;
+  uint8_t* mmapped_buffer_{ nullptr };
 };
 
 class MinuiBackendAdf : public MinuiBackend {
@@ -54,5 +63,3 @@ class MinuiBackendAdf : public MinuiBackend {
   unsigned int n_surfaces;
   GRSurfaceAdf surfaces[2];
 };
-
-#endif  // _GRAPHICS_ADF_H_
diff --git a/minui/graphics_drm.cpp b/minui/graphics_drm.cpp
index 630b80180..81b49fd95 100644
--- a/minui/graphics_drm.cpp
+++ b/minui/graphics_drm.cpp
@@ -70,8 +70,8 @@ void MinuiBackendDrm::Blank(bool blank) {
 void MinuiBackendDrm::DrmDestroySurface(GRSurfaceDrm* surface) {
   if (!surface) return;
 
-  if (surface->data) {
-    munmap(surface->data, surface->row_bytes * surface->height);
+  if (surface->mmapped_buffer_) {
+    munmap(surface->mmapped_buffer_, surface->row_bytes * surface->height);
   }
 
   if (surface->fb_id) {
@@ -172,15 +172,14 @@ GRSurfaceDrm* MinuiBackendDrm::DrmCreateSurface(int width, int height) {
   surface->width = width;
   surface->row_bytes = create_dumb.pitch;
   surface->pixel_bytes = create_dumb.bpp / 8;
-  surface->data = static_cast<unsigned char*>(mmap(nullptr, surface->height * surface->row_bytes,
-                                                   PROT_READ | PROT_WRITE, MAP_SHARED, drm_fd,
-                                                   map_dumb.offset));
-  if (surface->data == MAP_FAILED) {
+  auto mmapped = mmap(nullptr, surface->height * surface->row_bytes, PROT_READ | PROT_WRITE,
+                      MAP_SHARED, drm_fd, map_dumb.offset);
+  if (mmapped == MAP_FAILED) {
     perror("mmap() failed");
     DrmDestroySurface(surface);
     return nullptr;
   }
-
+  surface->mmapped_buffer_ = static_cast<uint8_t*>(mmapped);
   return surface;
 }
 
diff --git a/minui/graphics_drm.h b/minui/graphics_drm.h
index 756625b03..f3aad6bfc 100644
--- a/minui/graphics_drm.h
+++ b/minui/graphics_drm.h
@@ -14,8 +14,7 @@
  * limitations under the License.
  */
 
-#ifndef _GRAPHICS_DRM_H_
-#define _GRAPHICS_DRM_H_
+#pragma once
 
 #include <stdint.h>
 
@@ -25,11 +24,17 @@
 #include "minui/minui.h"
 
 class GRSurfaceDrm : public GRSurface {
+ public:
+  uint8_t* data() override {
+    return mmapped_buffer_;
+  }
+
  private:
+  friend class MinuiBackendDrm;
+
   uint32_t fb_id;
   uint32_t handle;
-
-  friend class MinuiBackendDrm;
+  uint8_t* mmapped_buffer_{ nullptr };
 };
 
 class MinuiBackendDrm : public MinuiBackend {
@@ -54,5 +59,3 @@ class MinuiBackendDrm : public MinuiBackend {
   drmModeConnector* main_monitor_connector;
   int drm_fd;
 };
-
-#endif  // _GRAPHICS_DRM_H_
diff --git a/minui/graphics_fbdev.cpp b/minui/graphics_fbdev.cpp
index 746f42aaa..f958d62d4 100644
--- a/minui/graphics_fbdev.cpp
+++ b/minui/graphics_fbdev.cpp
@@ -100,16 +100,16 @@ GRSurface* MinuiBackendFbdev::Init() {
   gr_framebuffer[0].height = vi.yres;
   gr_framebuffer[0].row_bytes = fi.line_length;
   gr_framebuffer[0].pixel_bytes = vi.bits_per_pixel / 8;
-  gr_framebuffer[0].data = static_cast<uint8_t*>(bits);
-  memset(gr_framebuffer[0].data, 0, gr_framebuffer[0].height * gr_framebuffer[0].row_bytes);
+  gr_framebuffer[0].buffer_ = static_cast<uint8_t*>(bits);
+  memset(gr_framebuffer[0].buffer_, 0, gr_framebuffer[0].height * gr_framebuffer[0].row_bytes);
 
   /* check if we can use double buffering */
   if (vi.yres * fi.line_length * 2 <= fi.smem_len) {
     double_buffered = true;
 
-    memcpy(gr_framebuffer + 1, gr_framebuffer, sizeof(GRSurface));
-    gr_framebuffer[1].data =
-        gr_framebuffer[0].data + gr_framebuffer[0].height * gr_framebuffer[0].row_bytes;
+    gr_framebuffer[1] = gr_framebuffer[0];
+    gr_framebuffer[1].buffer_ =
+        gr_framebuffer[0].buffer_ + gr_framebuffer[0].height * gr_framebuffer[0].row_bytes;
 
     gr_draw = gr_framebuffer + 1;
 
@@ -120,16 +120,12 @@ GRSurface* MinuiBackendFbdev::Init() {
     // draw in, and then "flipping" the buffer consists of a
     // memcpy from the buffer we allocated to the framebuffer.
 
-    gr_draw = static_cast<GRSurface*>(malloc(sizeof(GRSurface)));
-    memcpy(gr_draw, gr_framebuffer, sizeof(GRSurface));
-    gr_draw->data = static_cast<unsigned char*>(malloc(gr_draw->height * gr_draw->row_bytes));
-    if (!gr_draw->data) {
-      perror("failed to allocate in-memory surface");
-      return nullptr;
-    }
+    gr_draw = new GRSurfaceFbdev;
+    *gr_draw = gr_framebuffer[0];
+    gr_draw->buffer_ = new uint8_t[gr_draw->height * gr_draw->row_bytes];
   }
 
-  memset(gr_draw->data, 0, gr_draw->height * gr_draw->row_bytes);
+  memset(gr_draw->buffer_, 0, gr_draw->height * gr_draw->row_bytes);
   fb_fd = fd;
   SetDisplayedFramebuffer(0);
 
@@ -150,7 +146,7 @@ GRSurface* MinuiBackendFbdev::Flip() {
     SetDisplayedFramebuffer(1 - displayed_buffer);
   } else {
     // Copy from the in-memory surface to the framebuffer.
-    memcpy(gr_framebuffer[0].data, gr_draw->data, gr_draw->height * gr_draw->row_bytes);
+    memcpy(gr_framebuffer[0].buffer_, gr_draw->buffer_, gr_draw->height * gr_draw->row_bytes);
   }
   return gr_draw;
 }
@@ -160,8 +156,8 @@ MinuiBackendFbdev::~MinuiBackendFbdev() {
   fb_fd = -1;
 
   if (!double_buffered && gr_draw) {
-    free(gr_draw->data);
-    free(gr_draw);
+    delete[] gr_draw->buffer_;
+    delete gr_draw;
   }
   gr_draw = nullptr;
 }
diff --git a/minui/graphics_fbdev.h b/minui/graphics_fbdev.h
index 107e19567..be813dccb 100644
--- a/minui/graphics_fbdev.h
+++ b/minui/graphics_fbdev.h
@@ -14,14 +14,27 @@
  * limitations under the License.
  */
 
-#ifndef _GRAPHICS_FBDEV_H_
-#define _GRAPHICS_FBDEV_H_
+#pragma once
 
 #include <linux/fb.h>
+#include <stdint.h>
 
 #include "graphics.h"
 #include "minui/minui.h"
 
+class GRSurfaceFbdev : public GRSurface {
+ public:
+  uint8_t* data() override {
+    return buffer_;
+  }
+
+ private:
+  friend class MinuiBackendFbdev;
+
+  // Points to the start of the buffer: either the mmap'd framebuffer or one allocated in-memory.
+  uint8_t* buffer_;
+};
+
 class MinuiBackendFbdev : public MinuiBackend {
  public:
   GRSurface* Init() override;
@@ -33,12 +46,10 @@ class MinuiBackendFbdev : public MinuiBackend {
  private:
   void SetDisplayedFramebuffer(unsigned n);
 
-  GRSurface gr_framebuffer[2];
+  GRSurfaceFbdev gr_framebuffer[2];
   bool double_buffered;
-  GRSurface* gr_draw;
+  GRSurfaceFbdev* gr_draw;
   int displayed_buffer;
   fb_var_screeninfo vi;
   int fb_fd;
 };
-
-#endif  // _GRAPHICS_FBDEV_H_
diff --git a/minui/include/minui/minui.h b/minui/include/minui/minui.h
index fa13ecdff..e9bd1c4f1 100644
--- a/minui/include/minui/minui.h
+++ b/minui/include/minui/minui.h
@@ -14,12 +14,13 @@
  * limitations under the License.
  */
 
-#ifndef _MINUI_H_
-#define _MINUI_H_
+#pragma once
 
+#include <stdint.h>
 #include <sys/types.h>
 
 #include <functional>
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -27,12 +28,31 @@
 // Graphics.
 //
 
-struct GRSurface {
+class GRSurface {
+ public:
+  GRSurface() = default;
+  virtual ~GRSurface();
+
+  // Creates and returns a GRSurface instance for the given data_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(size_t data_size);
+
+  virtual uint8_t* data() {
+    return data_;
+  }
+
+  const uint8_t* data() const {
+    return const_cast<const uint8_t*>(const_cast<GRSurface*>(this)->data());
+  }
+
   int width;
   int height;
   int row_bytes;
   int pixel_bytes;
-  unsigned char* data;
+
+ private:
+  uint8_t* data_{ nullptr };
 };
 
 struct GRFont {
@@ -75,7 +95,7 @@ void gr_clear();
 void gr_color(unsigned char r, unsigned char g, unsigned char b, unsigned char a);
 void gr_fill(int x1, int y1, int x2, int y2);
 
-void gr_texticon(int x, int y, GRSurface* icon);
+void gr_texticon(int x, int y, const GRSurface* icon);
 
 const GRFont* gr_sys_font();
 int gr_init_font(const char* name, GRFont** dest);
@@ -85,7 +105,7 @@ int gr_measure(const GRFont* font, const char* s);
 // Returns -1 if font is nullptr.
 int gr_font_size(const GRFont* font, int* x, int* y);
 
-void gr_blit(GRSurface* source, int sx, int sy, int w, int h, int dx, int dy);
+void gr_blit(const GRSurface* source, int sx, int sy, int w, int h, int dx, int dy);
 unsigned int gr_get_width(const GRSurface* surface);
 unsigned int gr_get_height(const GRSurface* surface);
 
@@ -165,5 +185,3 @@ std::vector<std::string> get_locales_in_png(const std::string& png_name);
 // Free a surface allocated by any of the res_create_*_surface()
 // functions.
 void res_free_surface(GRSurface* surface);
-
-#endif
diff --git a/minui/resources.cpp b/minui/resources.cpp
index 477fbe2a2..c01c1868e 100644
--- a/minui/resources.cpp
+++ b/minui/resources.cpp
@@ -37,18 +37,23 @@
 
 #include "minui/minui.h"
 
-#define SURFACE_DATA_ALIGNMENT 8
-
 static std::string g_resource_dir{ "/res/images" };
 
-static GRSurface* malloc_surface(size_t data_size) {
-    size_t size = sizeof(GRSurface) + data_size + SURFACE_DATA_ALIGNMENT;
-    unsigned char* temp = static_cast<unsigned char*>(malloc(size));
-    if (temp == NULL) return NULL;
-    GRSurface* surface = reinterpret_cast<GRSurface*>(temp);
-    surface->data = temp + sizeof(GRSurface) +
-        (SURFACE_DATA_ALIGNMENT - (sizeof(GRSurface) % SURFACE_DATA_ALIGNMENT));
-    return surface;
+std::unique_ptr<GRSurface> GRSurface::Create(size_t data_size) {
+  static constexpr size_t kSurfaceDataAlignment = 8;
+  std::unique_ptr<GRSurface> result = std::make_unique<GRSurface>();
+  size_t aligned_size =
+      (data_size + kSurfaceDataAlignment - 1) / kSurfaceDataAlignment * kSurfaceDataAlignment;
+  result->data_ = static_cast<uint8_t*>(aligned_alloc(kSurfaceDataAlignment, aligned_size));
+  if (result->data_ == nullptr) return nullptr;
+  return result;
+}
+
+GRSurface::~GRSurface() {
+  if (data_ != nullptr) {
+    free(data_);
+    data_ = nullptr;
+  }
 }
 
 PngHandler::PngHandler(const std::string& name) {
@@ -133,18 +138,18 @@ PngHandler::~PngHandler() {
 // framebuffer pixel format; they need to be modified if the
 // framebuffer format changes (but nothing else should).
 
-// Allocate and return a GRSurface* sufficient for storing an image of
-// the indicated size in the framebuffer pixel format.
-static GRSurface* init_display_surface(png_uint_32 width, png_uint_32 height) {
-    GRSurface* surface = malloc_surface(width * height * 4);
-    if (surface == NULL) return NULL;
+// Allocates and returns a GRSurface* sufficient for storing an image of the indicated size in the
+// framebuffer pixel format.
+static std::unique_ptr<GRSurface> init_display_surface(png_uint_32 width, png_uint_32 height) {
+  std::unique_ptr<GRSurface> surface = GRSurface::Create(width * height * 4);
+  if (!surface) return nullptr;
 
-    surface->width = width;
-    surface->height = height;
-    surface->row_bytes = width * 4;
-    surface->pixel_bytes = 4;
+  surface->width = width;
+  surface->height = height;
+  surface->row_bytes = width * 4;
+  surface->pixel_bytes = 4;
 
-    return surface;
+  return surface;
 }
 
 // Copy 'input_row' to 'output_row', transforming it to the
@@ -202,7 +207,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();
 
-  GRSurface* surface = init_display_surface(width, height);
+  std::unique_ptr<GRSurface> surface = init_display_surface(width, height);
   if (!surface) {
     return -8;
   }
@@ -215,11 +220,11 @@ int res_create_display_surface(const char* name, GRSurface** pSurface) {
   for (png_uint_32 y = 0; y < height; ++y) {
     std::vector<unsigned char> p_row(width * 4);
     png_read_row(png_ptr, p_row.data(), nullptr);
-    transform_rgb_to_draw(p_row.data(), surface->data + y * surface->row_bytes,
+    transform_rgb_to_draw(p_row.data(), surface->data() + y * surface->row_bytes,
                           png_handler.channels(), width);
   }
 
-  *pSurface = surface;
+  *pSurface = surface.release();
 
   return 0;
 }
@@ -272,11 +277,12 @@ int res_create_multi_display_surface(const char* name, int* frames, int* fps,
     goto exit;
   }
   for (int i = 0; i < *frames; ++i) {
-    surface[i] = init_display_surface(width, height / *frames);
-    if (!surface[i]) {
+    auto created_surface = init_display_surface(width, height / *frames);
+    if (!created_surface) {
       result = -8;
       goto exit;
     }
+    surface[i] = created_surface.release();
   }
 
   if (gr_pixel_format() == PixelFormat::ABGR || gr_pixel_format() == PixelFormat::BGRA) {
@@ -287,7 +293,7 @@ int res_create_multi_display_surface(const char* name, int* frames, int* fps,
     std::vector<unsigned char> p_row(width * 4);
     png_read_row(png_ptr, p_row.data(), nullptr);
     int frame = y % *frames;
-    unsigned char* out_row = surface[frame]->data + (y / *frames) * surface[frame]->row_bytes;
+    unsigned char* out_row = surface[frame]->data() + (y / *frames) * surface[frame]->row_bytes;
     transform_rgb_to_draw(p_row.data(), out_row, png_handler.channels(), width);
   }
 
@@ -319,7 +325,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();
 
-  GRSurface* surface = malloc_surface(width * height);
+  std::unique_ptr<GRSurface> surface = GRSurface::Create(width * height);
   if (!surface) {
     return -8;
   }
@@ -334,11 +340,11 @@ int res_create_alpha_surface(const char* name, GRSurface** pSurface) {
   }
 
   for (png_uint_32 y = 0; y < height; ++y) {
-    unsigned char* p_row = surface->data + y * surface->row_bytes;
+    unsigned char* p_row = surface->data() + y * surface->row_bytes;
     png_read_row(png_ptr, p_row, nullptr);
   }
 
-  *pSurface = surface;
+  *pSurface = surface.release();
 
   return 0;
 }
@@ -429,7 +435,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);
 
-      GRSurface* surface = malloc_surface(w * h);
+      std::unique_ptr<GRSurface> surface = GRSurface::Create(w * h);
       if (!surface) {
         return -8;
       }
@@ -440,10 +446,10 @@ int res_create_localized_alpha_surface(const char* name,
 
       for (int i = 0; i < h; ++i, ++y) {
         png_read_row(png_ptr, row.data(), nullptr);
-        memcpy(surface->data + i * w, row.data(), w);
+        memcpy(surface->data() + i * w, row.data(), w);
       }
 
-      *pSurface = surface;
+      *pSurface = surface.release();
       break;
     }
 
diff --git a/recovery.cpp b/recovery.cpp
index 3ea282fc0..d7bc6fd2f 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -397,23 +397,22 @@ static bool wipe_data(Device* device) {
 
 static InstallResult prompt_and_wipe_data(Device* device) {
   // Use a single string and let ScreenRecoveryUI handles the wrapping.
-  std::vector<std::string> headers{
+  std::vector<std::string> wipe_data_menu_headers{
     "Can't load Android system. Your data may be corrupt. "
     "If you continue to get this message, you may need to "
     "perform a factory data reset and erase all user data "
     "stored on this device.",
   };
   // clang-format off
-  std::vector<std::string> items {
+  std::vector<std::string> wipe_data_menu_items {
     "Try again",
     "Factory data reset",
   };
   // clang-format on
   for (;;) {
-    size_t chosen_item = ui->ShowMenu(
-        headers, items, 0, true,
+    size_t chosen_item = ui->ShowPromptWipeDataMenu(
+        wipe_data_menu_headers, wipe_data_menu_items,
         std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
-
     // If ShowMenu() returned RecoveryUI::KeyError::INTERRUPTED, WaitKey() was interrupted.
     if (chosen_item == static_cast<size_t>(RecoveryUI::KeyError::INTERRUPTED)) {
       return INSTALL_KEY_INTERRUPTED;
@@ -421,6 +420,8 @@ static InstallResult prompt_and_wipe_data(Device* device) {
     if (chosen_item != 1) {
       return INSTALL_SUCCESS;  // Just reboot, no wipe; not a failure, user asked for it
     }
+
+    // TODO(xunchang) localize the confirmation texts also.
     if (ask_to_wipe_data(device)) {
       if (wipe_data(device)) {
         return INSTALL_SUCCESS;
diff --git a/screen_ui.cpp b/screen_ui.cpp
index 181d58ecf..c538815b5 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -197,12 +197,9 @@ int TextMenu::DrawItems(int x, int y, int screen_width, bool long_press) const {
   return offset;
 }
 
-GraphicMenu::GraphicMenu(size_t max_width, size_t max_height, GRSurface* graphic_headers,
-                         const std::vector<GRSurface*>& graphic_items, size_t initial_selection,
-                         const DrawInterface& draw_funcs)
+GraphicMenu::GraphicMenu(GRSurface* graphic_headers, const std::vector<GRSurface*>& graphic_items,
+                         size_t initial_selection, const DrawInterface& draw_funcs)
     : Menu(initial_selection, draw_funcs),
-      max_width_(max_width),
-      max_height_(max_height),
       graphic_headers_(graphic_headers),
       graphic_items_(graphic_items) {}
 
@@ -223,6 +220,7 @@ int GraphicMenu::Select(int sel) {
 }
 
 int GraphicMenu::DrawHeader(int x, int y) const {
+  draw_funcs_.SetColor(UIElement::HEADER);
   draw_funcs_.DrawTextIcon(x, y, graphic_headers_);
   return graphic_headers_->height;
 }
@@ -253,15 +251,16 @@ int GraphicMenu::DrawItems(int x, int y, int screen_width, bool long_press) cons
   return offset;
 }
 
-bool GraphicMenu::Validate() const {
+bool GraphicMenu::Validate(size_t max_width, size_t max_height, GRSurface* graphic_headers,
+                           const std::vector<GRSurface*>& graphic_items) {
   int offset = 0;
-  if (!ValidateGraphicSurface(offset, graphic_headers_)) {
+  if (!ValidateGraphicSurface(max_width, max_height, offset, graphic_headers)) {
     return false;
   }
-  offset += graphic_headers_->height;
+  offset += graphic_headers->height;
 
-  for (const auto& item : graphic_items_) {
-    if (!ValidateGraphicSurface(offset, item)) {
+  for (const auto& item : graphic_items) {
+    if (!ValidateGraphicSurface(max_width, max_height, offset, item)) {
       return false;
     }
     offset += item->height;
@@ -270,7 +269,8 @@ bool GraphicMenu::Validate() const {
   return true;
 }
 
-bool GraphicMenu::ValidateGraphicSurface(int y, const GRSurface* surface) const {
+bool GraphicMenu::ValidateGraphicSurface(size_t max_width, size_t max_height, int y,
+                                         const GRSurface* surface) {
   if (!surface) {
     fprintf(stderr, "Graphic surface can not be null");
     return false;
@@ -282,11 +282,11 @@ bool GraphicMenu::ValidateGraphicSurface(int y, const GRSurface* surface) const
     return false;
   }
 
-  if (surface->width > max_width_ || surface->height > max_height_ - y) {
+  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,"
             " max_height: %zu, vertical offset: %d\n",
-            surface->width, surface->height, max_width_, max_height_, y);
+            surface->width, surface->height, max_width, max_height, y);
     return false;
   }
 
@@ -697,7 +697,6 @@ void ScreenRecoveryUI::draw_menu_and_text_buffer_locked(
     const std::vector<std::string>& help_message) {
   int y = margin_height_;
   if (menu_) {
-    static constexpr int kMenuIndent = 4;
     int x = margin_width_ + kMenuIndent;
 
     SetColor(UIElement::INFO);
@@ -836,6 +835,16 @@ bool ScreenRecoveryUI::InitTextParams() {
   return true;
 }
 
+// TODO(xunchang) load localized text icons for the menu. (Init for screenRecoveryUI but
+// not wearRecoveryUI).
+bool ScreenRecoveryUI::LoadWipeDataMenuText() {
+  wipe_data_menu_header_text_ = nullptr;
+  factory_data_reset_text_ = nullptr;
+  try_again_text_ = nullptr;
+
+  return true;
+}
+
 bool ScreenRecoveryUI::Init(const std::string& locale) {
   RecoveryUI::Init(locale);
 
@@ -876,6 +885,8 @@ bool ScreenRecoveryUI::Init(const std::string& locale) {
   LoadLocalizedBitmap("no_command_text", &no_command_text);
   LoadLocalizedBitmap("error_text", &error_text);
 
+  LoadWipeDataMenuText();
+
   LoadAnimation();
 
   // Keep the progress bar updated, even when the process is otherwise busy.
@@ -1104,14 +1115,36 @@ void ScreenRecoveryUI::ShowFile(const std::string& filename) {
   text_row_ = old_text_row;
 }
 
-void ScreenRecoveryUI::StartMenu(const std::vector<std::string>& headers,
-                                 const std::vector<std::string>& items, size_t initial_selection) {
-  std::lock_guard<std::mutex> lg(updateMutex);
+std::unique_ptr<Menu> ScreenRecoveryUI::CreateMenu(GRSurface* graphic_header,
+                                                   const std::vector<GRSurface*>& graphic_items,
+                                                   const std::vector<std::string>& text_headers,
+                                                   const std::vector<std::string>& text_items,
+                                                   size_t initial_selection) const {
+  // horizontal unusable area: margin width + menu indent
+  size_t max_width = ScreenWidth() - margin_width_ - kMenuIndent;
+  // vertical unusable area: margin height + title lines + helper message + high light bar.
+  // It is safe to reserve more space.
+  size_t max_height = ScreenHeight() - margin_height_ - char_height_ * (title_lines_.size() + 3);
+  if (GraphicMenu::Validate(max_width, max_height, graphic_header, graphic_items)) {
+    return std::make_unique<GraphicMenu>(graphic_header, graphic_items, initial_selection, *this);
+  }
+
+  fprintf(stderr, "Failed to initialize graphic menu, falling back to use the text menu.\n");
+
+  return CreateMenu(text_headers, text_items, initial_selection);
+}
+
+std::unique_ptr<Menu> ScreenRecoveryUI::CreateMenu(const std::vector<std::string>& text_headers,
+                                                   const std::vector<std::string>& text_items,
+                                                   size_t initial_selection) const {
   if (text_rows_ > 0 && text_cols_ > 1) {
-    menu_ = std::make_unique<TextMenu>(scrollable_menu_, text_rows_, text_cols_ - 1, headers, items,
-                                       initial_selection, char_height_, *this);
-    update_screen_locked();
+    return std::make_unique<TextMenu>(scrollable_menu_, text_rows_, text_cols_ - 1, text_headers,
+                                      text_items, initial_selection, char_height_, *this);
   }
+
+  fprintf(stderr, "Failed to create text menu, text_rows %zu, text_cols %zu.\n", text_rows_,
+          text_cols_);
+  return nullptr;
 }
 
 int ScreenRecoveryUI::SelectMenu(int sel) {
@@ -1127,17 +1160,7 @@ int ScreenRecoveryUI::SelectMenu(int sel) {
   return sel;
 }
 
-void ScreenRecoveryUI::EndMenu() {
-  std::lock_guard<std::mutex> lg(updateMutex);
-  if (menu_) {
-    menu_.reset();
-    update_screen_locked();
-  }
-}
-
-size_t ScreenRecoveryUI::ShowMenu(const std::vector<std::string>& headers,
-                                  const std::vector<std::string>& items, size_t initial_selection,
-                                  bool menu_only,
+size_t ScreenRecoveryUI::ShowMenu(std::unique_ptr<Menu>&& menu, bool menu_only,
                                   const std::function<int(int, bool)>& key_handler) {
   // Throw away keys pressed previously, so user doesn't accidentally trigger menu items.
   FlushKeys();
@@ -1146,9 +1169,13 @@ size_t ScreenRecoveryUI::ShowMenu(const std::vector<std::string>& headers,
   // menu.
   if (IsKeyInterrupted()) return static_cast<size_t>(KeyError::INTERRUPTED);
 
-  StartMenu(headers, items, initial_selection);
+  CHECK(menu != nullptr);
+
+  // Starts and displays the menu
+  menu_ = std::move(menu);
+  Redraw();
 
-  int selected = initial_selection;
+  int selected = menu_->selection();
   int chosen_item = -1;
   while (chosen_item < 0) {
     int key = WaitKey();
@@ -1160,7 +1187,8 @@ size_t ScreenRecoveryUI::ShowMenu(const std::vector<std::string>& headers,
         continue;
       } else {
         LOG(INFO) << "Timed out waiting for key input; rebooting.";
-        EndMenu();
+        menu_.reset();
+        Redraw();
         return static_cast<size_t>(KeyError::TIMED_OUT);
       }
     }
@@ -1186,10 +1214,37 @@ size_t ScreenRecoveryUI::ShowMenu(const std::vector<std::string>& headers,
     }
   }
 
-  EndMenu();
+  menu_.reset();
+  Redraw();
+
   return chosen_item;
 }
 
+size_t ScreenRecoveryUI::ShowMenu(const std::vector<std::string>& headers,
+                                  const std::vector<std::string>& items, size_t initial_selection,
+                                  bool menu_only,
+                                  const std::function<int(int, bool)>& key_handler) {
+  auto menu = CreateMenu(headers, items, initial_selection);
+  if (menu == nullptr) {
+    return initial_selection;
+  }
+
+  return ShowMenu(CreateMenu(headers, items, initial_selection), menu_only, key_handler);
+}
+
+size_t ScreenRecoveryUI::ShowPromptWipeDataMenu(const std::vector<std::string>& backup_headers,
+                                                const std::vector<std::string>& backup_items,
+                                                const std::function<int(int, bool)>& key_handler) {
+  auto wipe_data_menu =
+      CreateMenu(wipe_data_menu_header_text_, { try_again_text_, factory_data_reset_text_ },
+                 backup_headers, backup_items, 0);
+  if (wipe_data_menu == nullptr) {
+    return 0;
+  }
+
+  return ShowMenu(std::move(wipe_data_menu), true, key_handler);
+}
+
 bool ScreenRecoveryUI::IsTextVisible() {
   std::lock_guard<std::mutex> lg(updateMutex);
   int visible = show_text;
diff --git a/screen_ui.h b/screen_ui.h
index 915288793..b1be100f9 100644
--- a/screen_ui.h
+++ b/screen_ui.h
@@ -29,7 +29,7 @@
 #include "ui.h"
 
 // From minui/minui.h.
-struct GRSurface;
+class GRSurface;
 
 enum class UIElement {
   HEADER,
@@ -167,25 +167,23 @@ class GraphicMenu : public Menu {
  public:
   // Constructs a Menu instance with the given |headers|, |items| and properties. Sets the initial
   // selection to |initial_selection|.
-  GraphicMenu(size_t max_width, size_t max_height, GRSurface* graphic_headers,
-              const std::vector<GRSurface*>& graphic_items, size_t initial_selection,
-              const DrawInterface& draw_funcs);
+  GraphicMenu(GRSurface* graphic_headers, const std::vector<GRSurface*>& graphic_items,
+              size_t initial_selection, const DrawInterface& draw_funcs);
 
   int Select(int sel) override;
   int DrawHeader(int x, int y) const override;
   int DrawItems(int x, int y, int screen_width, bool long_press) const override;
 
   // Checks if all the header and items are valid GRSurfaces; and that they can fit in the area
-  // defined by |max_width_| and |max_height_|.
-  bool Validate() const;
+  // defined by |max_width| and |max_height|.
+  static bool Validate(size_t max_width, size_t max_height, GRSurface* graphic_headers,
+                       const std::vector<GRSurface*>& graphic_items);
 
- private:
   // Returns true if |surface| fits on the screen with a vertical offset |y|.
-  bool ValidateGraphicSurface(int y, const GRSurface* surface) const;
-
-  const size_t max_width_;
-  const size_t max_height_;
+  static bool ValidateGraphicSurface(size_t max_width, size_t max_height, int y,
+                                     const GRSurface* surface);
 
+ private:
   // Pointers to the menu headers and items in graphic icons. This class does not have the ownership
   // of the these objects.
   GRSurface* graphic_headers_;
@@ -238,7 +236,13 @@ class ScreenRecoveryUI : public RecoveryUI, public DrawInterface {
   // the on-device resource files and shows the localized text, for manual inspection.
   void CheckBackgroundTextImages();
 
+  // Displays the localized wipe data menu.
+  size_t ShowPromptWipeDataMenu(const std::vector<std::string>& backup_headers,
+                                const std::vector<std::string>& backup_items,
+                                const std::function<int(int, bool)>& key_handler) override;
+
  protected:
+  static constexpr int kMenuIndent = 4;
   // The margin that we don't want to use for showing texts (e.g. round screen, or screen with
   // rounded corners).
   const int margin_width_;
@@ -252,18 +256,31 @@ class ScreenRecoveryUI : public RecoveryUI, public DrawInterface {
 
   virtual bool InitTextParams();
 
-  // Displays some header text followed by a menu of items, which appears at the top of the screen
-  // (in place of any scrolling ui_print() output, if necessary).
-  virtual void StartMenu(const std::vector<std::string>& headers,
-                         const std::vector<std::string>& items, size_t initial_selection);
+  virtual bool LoadWipeDataMenuText();
+
+  // Creates a GraphicMenu with |graphic_header| and |graphic_items|. If the GraphicMenu isn't
+  // valid or it doesn't fit on the screen; falls back to create a TextMenu instead. If succeeds,
+  // returns a unique pointer to the created menu; otherwise returns nullptr.
+  virtual std::unique_ptr<Menu> CreateMenu(GRSurface* graphic_header,
+                                           const std::vector<GRSurface*>& graphic_items,
+                                           const std::vector<std::string>& text_headers,
+                                           const std::vector<std::string>& text_items,
+                                           size_t initial_selection) const;
+
+  // Creates a TextMenu with |text_headers| and |text_items|; and sets the menu selection to
+  // |initial_selection|.
+  virtual std::unique_ptr<Menu> CreateMenu(const std::vector<std::string>& text_headers,
+                                           const std::vector<std::string>& text_items,
+                                           size_t initial_selection) const;
+
+  // Takes the ownership of |menu| and displays it.
+  virtual size_t ShowMenu(std::unique_ptr<Menu>&& menu, bool menu_only,
+                          const std::function<int(int, bool)>& key_handler);
 
   // Sets the menu highlight to the given index, wrapping if necessary. Returns the actual item
   // selected.
   virtual int SelectMenu(int sel);
 
-  // Ends menu mode, resetting the text overlay so that ui_print() statements will be displayed.
-  virtual void EndMenu();
-
   virtual void draw_background_locked();
   virtual void draw_foreground_locked();
   virtual void draw_screen_locked();
@@ -318,6 +335,11 @@ class ScreenRecoveryUI : public RecoveryUI, public DrawInterface {
   GRSurface* installing_text;
   GRSurface* no_command_text;
 
+  // Graphs for the wipe data menu
+  GRSurface* wipe_data_menu_header_text_;
+  GRSurface* try_again_text_;
+  GRSurface* factory_data_reset_text_;
+
   GRSurface** introFrames;
   GRSurface** loopFrames;
 
diff --git a/stub_ui.h b/stub_ui.h
index a3cf12b05..ca137dff3 100644
--- a/stub_ui.h
+++ b/stub_ui.h
@@ -68,6 +68,12 @@ class StubRecoveryUI : public RecoveryUI {
     return initial_selection;
   }
 
+  size_t ShowPromptWipeDataMenu(const std::vector<std::string>& /* backup_headers */,
+                                const std::vector<std::string>& /* backup_items */,
+                                const std::function<int(int, bool)>& /* key_handle */) override {
+    return 0;
+  }
+
   void SetTitle(const std::vector<std::string>& /* lines */) override {}
 };
 
diff --git a/tests/unit/minui_test.cpp b/tests/unit/minui_test.cpp
new file mode 100644
index 000000000..cad6a3d79
--- /dev/null
+++ b/tests/unit/minui_test.cpp
@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <stdint.h>
+
+#include <memory>
+
+#include <gtest/gtest.h>
+
+#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++) {
+    std::unique_ptr<GRSurface> surface(GRSurface::Create(data_size));
+    ASSERT_TRUE(surface);
+    ASSERT_EQ(0, reinterpret_cast<uintptr_t>(surface->data()) % kSurfaceDataAlignment);
+  }
+}
diff --git a/tests/unit/screen_ui_test.cpp b/tests/unit/screen_ui_test.cpp
index ec269503e..06a98c742 100644
--- a/tests/unit/screen_ui_test.cpp
+++ b/tests/unit/screen_ui_test.cpp
@@ -69,6 +69,17 @@ class ScreenUITest : public testing::Test {
   MockDrawFunctions draw_funcs_;
 };
 
+// TODO(xunchang) Create a constructor.
+static GRSurface CreateFakeGRSurface(int width, int height, int row_bytes, int pixel_bytes) {
+  GRSurface fake_surface;
+  fake_surface.width = width;
+  fake_surface.height = height;
+  fake_surface.row_bytes = row_bytes;
+  fake_surface.pixel_bytes = pixel_bytes;
+
+  return fake_surface;
+}
+
 TEST_F(ScreenUITest, StartPhoneMenuSmoke) {
   TextMenu menu(false, 10, 20, HEADERS, ITEMS, 0, 20, draw_funcs_);
   ASSERT_FALSE(menu.scrollable());
@@ -229,6 +240,43 @@ TEST_F(ScreenUITest, WearMenuSelectItemsOverflow) {
   ASSERT_EQ(3u, menu.MenuEnd());
 }
 
+TEST_F(ScreenUITest, GraphicMenuSelection) {
+  auto fake_surface = CreateFakeGRSurface(50, 50, 50, 1);
+  std::vector<GRSurface*> items = { &fake_surface, &fake_surface, &fake_surface };
+  GraphicMenu menu(&fake_surface, items, 0, draw_funcs_);
+
+  ASSERT_EQ(0, menu.selection());
+
+  int sel = 0;
+  for (int i = 0; i < 3; i++) {
+    sel = menu.Select(++sel);
+    ASSERT_EQ((i + 1) % 3, sel);
+    ASSERT_EQ(sel, menu.selection());
+  }
+
+  sel = 0;
+  for (int i = 0; i < 3; i++) {
+    sel = menu.Select(--sel);
+    ASSERT_EQ(2 - i, sel);
+    ASSERT_EQ(sel, menu.selection());
+  }
+}
+
+TEST_F(ScreenUITest, GraphicMenuValidate) {
+  auto fake_surface = CreateFakeGRSurface(50, 50, 50, 1);
+  std::vector<GRSurface*> items = { &fake_surface, &fake_surface, &fake_surface };
+
+  ASSERT_TRUE(GraphicMenu::Validate(200, 200, &fake_surface, items));
+
+  // Menu exceeds the horizontal boundary.
+  auto wide_surface = CreateFakeGRSurface(300, 50, 300, 1);
+  ASSERT_FALSE(GraphicMenu::Validate(299, 200, &wide_surface, items));
+
+  // Menu exceeds the vertical boundary.
+  items.push_back(&fake_surface);
+  ASSERT_FALSE(GraphicMenu::Validate(200, 249, &fake_surface, items));
+}
+
 static constexpr int kMagicAction = 101;
 
 enum class KeyCode : int {
diff --git a/ui.h b/ui.h
index e0fb13e40..1e6186a11 100644
--- a/ui.h
+++ b/ui.h
@@ -162,6 +162,13 @@ class RecoveryUI {
                           const std::vector<std::string>& items, size_t initial_selection,
                           bool menu_only, const std::function<int(int, bool)>& key_handler) = 0;
 
+  // Displays the localized wipe data menu with pre-generated graphs. If there's an issue
+  // with the graphs, falls back to use the backup string headers and items instead. The initial
+  // selection is the 0th item in the menu, which is expected to reboot the device without a wipe.
+  virtual size_t ShowPromptWipeDataMenu(const std::vector<std::string>& backup_headers,
+                                        const std::vector<std::string>& backup_items,
+                                        const std::function<int(int, bool)>& key_handler) = 0;
+
   // Resets the key interrupt status.
   void ResetKeyInterruptStatus() {
     key_interrupted_ = false;
diff --git a/wear_ui.cpp b/wear_ui.cpp
index 8f3bc7bbe..0611f94c9 100644
--- a/wear_ui.cpp
+++ b/wear_ui.cpp
@@ -95,13 +95,14 @@ void WearRecoveryUI::update_progress_locked() {
 
 void WearRecoveryUI::SetStage(int /* current */, int /* max */) {}
 
-void WearRecoveryUI::StartMenu(const std::vector<std::string>& headers,
-                               const std::vector<std::string>& items, size_t initial_selection) {
-  std::lock_guard<std::mutex> lg(updateMutex);
+std::unique_ptr<Menu> WearRecoveryUI::CreateMenu(const std::vector<std::string>& text_headers,
+                                                 const std::vector<std::string>& text_items,
+                                                 size_t initial_selection) const {
   if (text_rows_ > 0 && text_cols_ > 0) {
-    menu_ = std::make_unique<TextMenu>(scrollable_menu_, text_rows_ - menu_unusable_rows_ - 1,
-                                       text_cols_ - 1, headers, items, initial_selection,
-                                       char_height_, *this);
-    update_screen_locked();
+    return std::make_unique<TextMenu>(scrollable_menu_, text_rows_ - menu_unusable_rows_ - 1,
+                                      text_cols_ - 1, text_headers, text_items, initial_selection,
+                                      char_height_, *this);
   }
+
+  return nullptr;
 }
diff --git a/wear_ui.h b/wear_ui.h
index b80cfd758..429af69d2 100644
--- a/wear_ui.h
+++ b/wear_ui.h
@@ -36,8 +36,9 @@ class WearRecoveryUI : public ScreenRecoveryUI {
   // Recovery, build id and etc) and the bottom lines that may otherwise go out of the screen.
   const int menu_unusable_rows_;
 
-  void StartMenu(const std::vector<std::string>& headers, const std::vector<std::string>& items,
-                 size_t initial_selection) override;
+  std::unique_ptr<Menu> CreateMenu(const std::vector<std::string>& text_headers,
+                                   const std::vector<std::string>& text_items,
+                                   size_t initial_selection) const override;
 
   int GetProgressBaseline() const override;