From 9f80b9dbdc6540d35ce07a7c4a8058af852f7902 Mon Sep 17 00:00:00 2001
From: DianaOu <dianaou@google.com>
Date: Wed, 01 Jul 2026 13:48:40 -0700
Subject: [PATCH] [Lens Overlay] Add guards to SaveAsImage Mojo handler

This CL adds tab activation and initialization data guards to
LensOverlayController::SaveAsImage to prevent background execution and
potential null pointer dereferences.

Bug: b:522396262
Change-Id: Ie48f885804ec0c7c29d38e88d50faef2b1019531
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/8031806
Reviewed-by: Juan Mojica <juanmojica@google.com>
Commit-Queue: Juan Mojica <juanmojica@google.com>
Auto-Submit: Diana OuYang <dianaou@google.com>
Cr-Commit-Position: refs/heads/main@{#1655734}
---

diff --git a/chrome/browser/ui/lens/BUILD.gn b/chrome/browser/ui/lens/BUILD.gn
index b8dd73e..229198ec 100644
--- a/chrome/browser/ui/lens/BUILD.gn
+++ b/chrome/browser/ui/lens/BUILD.gn
@@ -345,6 +345,7 @@
     "//ui/compositor",
     "//ui/events:events_base",
     "//ui/events:test_support",
+    "//ui/shell_dialogs:test_support",
     "//ui/views",
     "//ui/views:test_support",
     "//ui/views/controls/webview",
diff --git a/chrome/browser/ui/lens/lens_overlay_controller.cc b/chrome/browser/ui/lens/lens_overlay_controller.cc
index 49aa354..41c35f9 100644
--- a/chrome/browser/ui/lens/lens_overlay_controller.cc
+++ b/chrome/browser/ui/lens/lens_overlay_controller.cc
@@ -536,6 +536,13 @@
 
 void LensOverlayController::SaveAsImage(
     lens::mojom::CenterRotatedBoxPtr region) {
+  if (!tab_->IsActivated()) {
+    return;
+  }
+  if (!initialization_data_ ||
+      initialization_data_->initial_screenshot_.drawsNothing()) {
+    return;
+  }
   SkBitmap cropped = lens::CropBitmapToRegion(
       initialization_data_->initial_screenshot_, std::move(region));
   const GURL data_url = GURL(webui::GetBitmapDataUrl(cropped));
diff --git a/chrome/browser/ui/lens/lens_overlay_controller_browsertest.cc b/chrome/browser/ui/lens/lens_overlay_controller_browsertest.cc
index 3ab52ca..802047f 100644
--- a/chrome/browser/ui/lens/lens_overlay_controller_browsertest.cc
+++ b/chrome/browser/ui/lens/lens_overlay_controller_browsertest.cc
@@ -28,6 +28,7 @@
 #include "base/test/scoped_feature_list.h"
 #include "base/test/with_feature_override.h"
 #include "base/threading/thread_restrictions.h"
+#include "base/task/single_thread_task_runner.h"
 #include "build/build_config.h"
 #include "chrome/browser/autocomplete/aim_eligibility_service_factory.h"
 #include "chrome/browser/companion/text_finder/text_highlighter.h"
@@ -162,6 +163,7 @@
 #include "ui/events/base_event_utils.h"
 #include "ui/events/event_constants.h"
 #include "ui/events/test/test_event.h"
+#include "ui/shell_dialogs/fake_select_file_dialog.h"
 #include "ui/views/accessibility/view_accessibility.h"
 #include "ui/views/controls/webview/webview.h"
 #include "ui/views/interaction/element_tracker_views.h"
@@ -9939,3 +9941,75 @@
       clipboard, ui::ClipboardFormatType::BitmapType(),
       ui::ClipboardBuffer::kCopyPaste, /* data_dst = */ nullptr));
 }
+
+IN_PROC_BROWSER_TEST_F(LensOverlayControllerBrowserTest,
+                       SaveAsImageBackgroundCheck) {
+  WaitForPaint();
+
+  auto* controller = GetLensOverlayController();
+  ASSERT_EQ(controller->state(), State::kOff);
+
+  // Show the overlay.
+  OpenLensOverlay(LensOverlayInvocationSource::kAppMenu);
+  ASSERT_EQ(controller->state(), State::kScreenshot);
+  ASSERT_TRUE(base::test::RunUntil(
+      [&]() { return controller->state() == State::kOverlay; }));
+
+  // Keep track of the active tab index.
+  int active_controller_tab_index =
+      browser()->tab_strip_model()->active_index();
+
+  // Background the tab by opening a new tab.
+  WaitForPaint(kDocumentWithNamedElement,
+               WindowOpenDisposition::NEW_FOREGROUND_TAB,
+               ui_test_utils::BROWSER_TEST_WAIT_FOR_TAB |
+                   ui_test_utils::BROWSER_TEST_WAIT_FOR_LOAD_STOP);
+  EXPECT_TRUE(base::test::RunUntil(
+      [&]() { return controller->state() == State::kBackground; }));
+
+  // Register the fake file dialog factory.
+  ui::FakeSelectFileDialog::Factory* factory =
+      ui::FakeSelectFileDialog::RegisterFactory();
+  bool file_dialog_opened = false;
+  factory->SetOpenCallback(base::BindRepeating(
+      [](bool* opened) { *opened = true; }, &file_dialog_opened));
+
+  // Test SaveAsImage when backgrounded. It should NOT open the dialog.
+  lens::mojom::LensPageHandler* page_handler = controller;
+  auto region = lens::mojom::CenterRotatedBox::New();
+  region->box = gfx::RectF(0.1, 0.1, 0.2, 0.2);
+  region->coordinate_type =
+      lens::mojom::CenterRotatedBox::CoordinateType::kNormalized;
+
+  page_handler->SaveAsImage(std::move(region));
+
+  // Yield to the message loop to ensure any pending tasks are run.
+  base::RunLoop run_loop;
+  base::SingleThreadTaskRunner::GetCurrentDefault()->PostTask(
+      FROM_HERE, run_loop.QuitClosure());
+  run_loop.Run();
+
+  // Dialog should NOT have been opened.
+  EXPECT_FALSE(file_dialog_opened);
+
+  // Reactivate the tab.
+  browser()->tab_strip_model()->ActivateTabAt(active_controller_tab_index);
+  EXPECT_TRUE(base::test::RunUntil(
+      [&]() { return controller->state() == State::kOverlay; }));
+
+  // Test SaveAsImage when active. It should open the dialog.
+  auto region2 = lens::mojom::CenterRotatedBox::New();
+  region2->box = gfx::RectF(0.1, 0.1, 0.2, 0.2);
+  region2->coordinate_type =
+      lens::mojom::CenterRotatedBox::CoordinateType::kNormalized;
+
+  page_handler->SaveAsImage(std::move(region2));
+
+  // Wait for the dialog to open.
+  ASSERT_TRUE(base::test::RunUntil([&]() { return file_dialog_opened; }));
+
+  // Complete the dialog to avoid hanging download manager.
+  auto* dialog = factory->GetLastDialog();
+  ASSERT_TRUE(dialog);
+  dialog->CallFileSelectionCanceled();
+}

diff --git a/chrome/browser/ui/lens/lens_overlay_controller_browsertest.cc b/chrome/browser/ui/lens/lens_overlay_controller_browsertest.cc
index 3ab52ca..802047f 100644
--- a/chrome/browser/ui/lens/lens_overlay_controller_browsertest.cc
+++ b/chrome/browser/ui/lens/lens_overlay_controller_browsertest.cc
@@ -28,6 +28,7 @@
 #include "base/test/scoped_feature_list.h"
 #include "base/test/with_feature_override.h"
 #include "base/threading/thread_restrictions.h"
+#include "base/task/single_thread_task_runner.h"
 #include "build/build_config.h"
 #include "chrome/browser/autocomplete/aim_eligibility_service_factory.h"
 #include "chrome/browser/companion/text_finder/text_highlighter.h"
@@ -162,6 +163,7 @@
 #include "ui/events/base_event_utils.h"
 #include "ui/events/event_constants.h"
 #include "ui/events/test/test_event.h"
+#include "ui/shell_dialogs/fake_select_file_dialog.h"
 #include "ui/views/accessibility/view_accessibility.h"
 #include "ui/views/controls/webview/webview.h"
 #include "ui/views/interaction/element_tracker_views.h"
@@ -9939,3 +9941,75 @@
       clipboard, ui::ClipboardFormatType::BitmapType(),
       ui::ClipboardBuffer::kCopyPaste, /* data_dst = */ nullptr));
 }
+
+IN_PROC_BROWSER_TEST_F(LensOverlayControllerBrowserTest,
+                       SaveAsImageBackgroundCheck) {
+  WaitForPaint();
+
+  auto* controller = GetLensOverlayController();
+  ASSERT_EQ(controller->state(), State::kOff);
+
+  // Show the overlay.
+  OpenLensOverlay(LensOverlayInvocationSource::kAppMenu);
+  ASSERT_EQ(controller->state(), State::kScreenshot);
+  ASSERT_TRUE(base::test::RunUntil(
+      [&]() { return controller->state() == State::kOverlay; }));
+
+  // Keep track of the active tab index.
+  int active_controller_tab_index =
+      browser()->tab_strip_model()->active_index();
+
+  // Background the tab by opening a new tab.
+  WaitForPaint(kDocumentWithNamedElement,
+               WindowOpenDisposition::NEW_FOREGROUND_TAB,
+               ui_test_utils::BROWSER_TEST_WAIT_FOR_TAB |
+                   ui_test_utils::BROWSER_TEST_WAIT_FOR_LOAD_STOP);
+  EXPECT_TRUE(base::test::RunUntil(
+      [&]() { return controller->state() == State::kBackground; }));
+
+  // Register the fake file dialog factory.
+  ui::FakeSelectFileDialog::Factory* factory =
+      ui::FakeSelectFileDialog::RegisterFactory();
+  bool file_dialog_opened = false;
+  factory->SetOpenCallback(base::BindRepeating(
+      [](bool* opened) { *opened = true; }, &file_dialog_opened));
+
+  // Test SaveAsImage when backgrounded. It should NOT open the dialog.
+  lens::mojom::LensPageHandler* page_handler = controller;
+  auto region = lens::mojom::CenterRotatedBox::New();
+  region->box = gfx::RectF(0.1, 0.1, 0.2, 0.2);
+  region->coordinate_type =
+      lens::mojom::CenterRotatedBox::CoordinateType::kNormalized;
+
+  page_handler->SaveAsImage(std::move(region));
+
+  // Yield to the message loop to ensure any pending tasks are run.
+  base::RunLoop run_loop;
+  base::SingleThreadTaskRunner::GetCurrentDefault()->PostTask(
+      FROM_HERE, run_loop.QuitClosure());
+  run_loop.Run();
+
+  // Dialog should NOT have been opened.
+  EXPECT_FALSE(file_dialog_opened);
+
+  // Reactivate the tab.
+  browser()->tab_strip_model()->ActivateTabAt(active_controller_tab_index);
+  EXPECT_TRUE(base::test::RunUntil(
+      [&]() { return controller->state() == State::kOverlay; }));
+
+  // Test SaveAsImage when active. It should open the dialog.
+  auto region2 = lens::mojom::CenterRotatedBox::New();
+  region2->box = gfx::RectF(0.1, 0.1, 0.2, 0.2);
+  region2->coordinate_type =
+      lens::mojom::CenterRotatedBox::CoordinateType::kNormalized;
+
+  page_handler->SaveAsImage(std::move(region2));
+
+  // Wait for the dialog to open.
+  ASSERT_TRUE(base::test::RunUntil([&]() { return file_dialog_opened; }));
+
+  // Complete the dialog to avoid hanging download manager.
+  auto* dialog = factory->GetLastDialog();
+  ASSERT_TRUE(dialog);
+  dialog->CallFileSelectionCanceled();
+}