DLP: Enforce screenshot restrictions in ChromeScreenshotGrabber.

Full- and partial-screenshots are restricted if any restricted for screenshots data is visible.
Window-screenshots are restricted if the window contains any restricted for screenshots WebContents.

Bug: 1109723
Test: Manual, new unit/browser tests being added.
Change-Id: I62bd96f39706b13ef4468d78b0ce9258699c985f
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2317961Reviewed-by: Scott Violet <sky@chromium.org>
Reviewed-by: Nikita Podguzov <nikitapodguzov@chromium.org>
Commit-Queue: Sergey Poromov <poromov@chromium.org>
Cr-Commit-Position: refs/heads/master@{#796229}

chrome/browser/chromeos/BUILD.gn

@@ -3306,6 +3306,8 @@ source_set("unit_tests") {
     "policy/device_policy_decoder_chromeos_unittest.cc",
     "policy/dlp/dlp_content_manager_unittest.cc",
     "policy/dlp/dlp_content_tab_helper_unittest.cc",
+    "policy/dlp/mock_dlp_content_manager.cc",
+    "policy/dlp/mock_dlp_content_manager.h",
     "policy/dm_token_storage_unittest.cc",
     "policy/extension_cache_unittest.cc",
     "policy/extension_install_event_log_collector_unittest.cc",

chrome/browser/chromeos/policy/dlp/dlp_content_manager.cc

@@ -8,8 +8,12 @@
 #include "base/bind.h"
 #include "base/stl_util.h"
 #include "base/threading/thread_task_runner_handle.h"
+#include "chrome/browser/ui/ash/chrome_screenshot_grabber.h"
 #include "content/public/browser/visibility.h"
 #include "content/public/browser/web_contents.h"
+#include "ui/aura/window.h"
+#include "ui/gfx/geometry/rect.h"
+#include "ui/gfx/skia_util.h"
 #include "url/gurl.h"
 
 namespace policy {
@@ -32,7 +36,7 @@ DlpContentManager* DlpContentManager::Get() {
 }
 
 DlpContentRestrictionSet DlpContentManager::GetConfidentialRestrictions(
-    const content::WebContents* web_contents) const {
+    content::WebContents* web_contents) const {
   if (!base::Contains(confidential_web_contents_, web_contents))
     return DlpContentRestrictionSet();
   return confidential_web_contents_.at(web_contents);
@@ -43,6 +47,57 @@ DlpContentRestrictionSet DlpContentManager::GetOnScreenPresentRestrictions()
   return on_screen_restrictions_;
 }
 
+bool DlpContentManager::IsScreenshotRestricted(
+    const ScreenshotArea& area) const {
+  // Fullscreen - restricted if any confidential data is visible.
+  if (area.type == ScreenshotType::kAllRootWindows) {
+    return GetOnScreenPresentRestrictions().HasRestriction(
+        DlpContentRestriction::kScreenshot);
+  }
+
+  // Window - restricted if the window contains confidential data.
+  if (area.type == ScreenshotType::kWindow) {
+    DCHECK(area.window);
+    for (auto& entry : confidential_web_contents_) {
+      aura::Window* web_contents_window = entry.first->GetNativeView();
+      if (entry.second.HasRestriction(DlpContentRestriction::kScreenshot) &&
+          area.window->Contains(web_contents_window)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  DCHECK_EQ(area.type, ScreenshotType::kPartialWindow);
+  DCHECK(area.rect);
+  DCHECK(area.window);
+  // Partial - restricted if any visible confidential WebContents intersects
+  // with the area.
+  for (auto& entry : confidential_web_contents_) {
+    if (entry.first->GetVisibility() != content::Visibility::VISIBLE ||
+        !entry.second.HasRestriction(DlpContentRestriction::kScreenshot)) {
+      continue;
+    }
+    aura::Window* web_contents_window = entry.first->GetNativeView();
+    aura::Window* root_window = web_contents_window->GetRootWindow();
+    // If no root window, then the WebContent shouldn't be visible.
+    if (!root_window)
+      continue;
+    // Not allowing if the area intersects with confidential WebContents,
+    // but the intersection doesn't belong to occluded area.
+    gfx::Rect intersection(*area.rect);
+    aura::Window::ConvertRectToTarget(area.window, root_window, &intersection);
+    intersection.Intersect(web_contents_window->GetBoundsInRootWindow());
+    if (!intersection.IsEmpty() &&
+        !web_contents_window->occluded_region().contains(
+            gfx::RectToSkIRect(intersection))) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
 /* static */
 void DlpContentManager::SetDlpContentManagerForTesting(
     DlpContentManager* dlp_content_manager) {
@@ -74,7 +129,7 @@ void DlpContentManager::OnConfidentialityChanged(
 }
 
 void DlpContentManager::OnWebContentsDestroyed(
-    const content::WebContents* web_contents) {
+    content::WebContents* web_contents) {
   RemoveFromConfidential(web_contents);
 }
 
@@ -91,7 +146,7 @@ void DlpContentManager::OnVisibilityChanged(
 }
 
 void DlpContentManager::RemoveFromConfidential(
-    const content::WebContents* web_contents) {
+    content::WebContents* web_contents) {
   confidential_web_contents_.erase(web_contents);
   MaybeChangeOnScreenRestrictions();
 }

chrome/browser/chromeos/policy/dlp/dlp_content_manager.h

@@ -6,10 +6,12 @@
 #define CHROME_BROWSER_CHROMEOS_POLICY_DLP_DLP_CONTENT_MANAGER_H_
 
 #include "base/containers/flat_map.h"
+#include "base/gtest_prod_util.h"
 #include "base/time/time.h"
 #include "chrome/browser/chromeos/policy/dlp/dlp_content_restriction_set.h"
 
 class GURL;
+struct ScreenshotArea;
 
 namespace content {
 class WebContents;
@@ -30,12 +32,15 @@ class DlpContentManager {
   // Returns which restrictions are applied to the |web_contents| according to
   // the policy.
   DlpContentRestrictionSet GetConfidentialRestrictions(
-      const content::WebContents* web_contents) const;
+      content::WebContents* web_contents) const;
 
   // Returns which restrictions are applied to the WebContents which are
   // currently visible.
   DlpContentRestrictionSet GetOnScreenPresentRestrictions() const;
 
+  // Returns whether screenshots should be restricted.
+  virtual bool IsScreenshotRestricted(const ScreenshotArea& area) const;
+
   // The caller (test) should manage |dlp_content_manager| lifetime.
   // Reset doesn't delete the object.
   static void SetDlpContentManagerForTesting(
@@ -43,6 +48,7 @@ class DlpContentManager {
   static void ResetDlpContentManagerForTesting();
 
  private:
+  FRIEND_TEST_ALL_PREFIXES(DlpContentManagerBrowserTest, ScreenshotsRestricted);
   friend class DlpContentManagerTest;
   friend class DlpContentTabHelper;
   friend class MockDlpContentManager;
@@ -59,7 +65,7 @@ class DlpContentManager {
       content::WebContents* web_contents,
       const DlpContentRestrictionSet& restriction_set);
   // Called when |web_contents| is about to be destroyed.
-  virtual void OnWebContentsDestroyed(const content::WebContents* web_contents);
+  virtual void OnWebContentsDestroyed(content::WebContents* web_contents);
   // Should return which restrictions are being applied to the |url| according
   // to the policies.
   virtual DlpContentRestrictionSet GetRestrictionSetForURL(
@@ -68,7 +74,7 @@ class DlpContentManager {
   virtual void OnVisibilityChanged(content::WebContents* web_contents);
 
   // Helper to remove |web_contents| from the confidential set.
-  void RemoveFromConfidential(const content::WebContents* web_contents);
+  void RemoveFromConfidential(content::WebContents* web_contents);
 
   // Updates |on_screen_restrictions_| and calls
   // OnScreenRestrictionsChanged() if needed.

chrome/browser/chromeos/policy/dlp/dlp_content_manager_browsertest.cc

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chrome/browser/chromeos/policy/dlp/dlp_content_manager.h"
+
+#include "chrome/browser/ui/ash/chrome_screenshot_grabber.h"
+#include "chrome/browser/ui/browser.h"
+#include "chrome/browser/ui/browser_window.h"
+#include "chrome/test/base/in_process_browser_test.h"
+#include "chrome/test/base/ui_test_utils.h"
+#include "content/public/test/browser_test.h"
+#include "testing/gtest/include/gtest/gtest.h"
+#include "ui/aura/window.h"
+#include "ui/gfx/geometry/rect.h"
+
+namespace policy {
+
+namespace {
+const DlpContentRestrictionSet kScreenshotRestricted(
+    DlpContentRestriction::kScreenshot);
+}  // namespace
+
+class DlpContentManagerBrowserTest : public InProcessBrowserTest {
+ public:
+  DlpContentManagerBrowserTest() {}
+};
+
+IN_PROC_BROWSER_TEST_F(DlpContentManagerBrowserTest, ScreenshotsRestricted) {
+  DlpContentManager* manager = DlpContentManager::Get();
+  ui_test_utils::NavigateToURL(browser(), GURL("https://example.com"));
+  content::WebContents* web_contents =
+      browser()->tab_strip_model()->GetActiveWebContents();
+  aura::Window* root_window =
+      browser()->window()->GetNativeWindow()->GetRootWindow();
+  ScreenshotArea fullscreen = ScreenshotArea::CreateForAllRootWindows();
+  ScreenshotArea window = ScreenshotArea::CreateForWindow(root_window);
+  const gfx::Rect web_contents_rect = web_contents->GetContainerBounds();
+  gfx::Rect out_rect(web_contents_rect);
+  out_rect.Offset(web_contents_rect.width(), web_contents_rect.height());
+  gfx::Rect in_rect(web_contents_rect);
+  in_rect.Offset(web_contents_rect.width() / 2, web_contents_rect.height() / 2);
+  ScreenshotArea partial_out =
+      ScreenshotArea::CreateForPartialWindow(root_window, out_rect);
+  ScreenshotArea partial_in =
+      ScreenshotArea::CreateForPartialWindow(root_window, in_rect);
+
+  EXPECT_FALSE(manager->IsScreenshotRestricted(fullscreen));
+  EXPECT_FALSE(manager->IsScreenshotRestricted(window));
+  EXPECT_FALSE(manager->IsScreenshotRestricted(partial_in));
+  EXPECT_FALSE(manager->IsScreenshotRestricted(partial_out));
+
+  manager->OnConfidentialityChanged(web_contents, kScreenshotRestricted);
+  EXPECT_TRUE(manager->IsScreenshotRestricted(fullscreen));
+  EXPECT_TRUE(manager->IsScreenshotRestricted(window));
+  EXPECT_TRUE(manager->IsScreenshotRestricted(partial_in));
+  EXPECT_FALSE(manager->IsScreenshotRestricted(partial_out));
+
+  web_contents->WasHidden();
+  manager->OnVisibilityChanged(web_contents);
+  EXPECT_FALSE(manager->IsScreenshotRestricted(fullscreen));
+  EXPECT_TRUE(manager->IsScreenshotRestricted(window));
+  EXPECT_FALSE(manager->IsScreenshotRestricted(partial_in));
+  EXPECT_FALSE(manager->IsScreenshotRestricted(partial_out));
+
+  web_contents->WasShown();
+  manager->OnVisibilityChanged(web_contents);
+  EXPECT_TRUE(manager->IsScreenshotRestricted(fullscreen));
+  EXPECT_TRUE(manager->IsScreenshotRestricted(window));
+  EXPECT_TRUE(manager->IsScreenshotRestricted(partial_in));
+  EXPECT_FALSE(manager->IsScreenshotRestricted(partial_out));
+
+  manager->OnWebContentsDestroyed(web_contents);
+  EXPECT_FALSE(manager->IsScreenshotRestricted(fullscreen));
+  EXPECT_FALSE(manager->IsScreenshotRestricted(partial_in));
+  EXPECT_FALSE(manager->IsScreenshotRestricted(partial_out));
+}
+
+}  // namespace policy

chrome/browser/chromeos/policy/dlp/dlp_content_manager_unittest.cc

@@ -17,8 +17,9 @@ namespace policy {
 
 namespace {
 const DlpContentRestrictionSet kEmptyRestrictionSet;
-const DlpContentRestrictionSet kNonEmptyRestrictionSet(
+const DlpContentRestrictionSet kScreenshotRestricted(
     DlpContentRestriction::kScreenshot);
+const DlpContentRestrictionSet kNonEmptyRestrictionSet = kScreenshotRestricted;
 const DlpContentRestrictionSet kPrivacyScreenEnforced(
     DlpContentRestriction::kPrivacyScreen);
 

chrome/browser/chromeos/policy/dlp/dlp_content_tab_helper_unittest.cc

@@ -4,7 +4,7 @@
 
 #include "chrome/browser/chromeos/policy/dlp/dlp_content_tab_helper.h"
 
-#include "chrome/browser/chromeos/policy/dlp/dlp_content_manager.h"
+#include "chrome/browser/chromeos/policy/dlp/mock_dlp_content_manager.h"
 #include "chrome/browser/ui/browser.h"
 #include "chrome/browser/ui/tabs/tab_activity_simulator.h"
 #include "chrome/test/base/chrome_render_view_host_test_harness.h"
@@ -25,16 +25,6 @@ const DlpContentRestrictionSet kNonEmptyRestrictionSet(
     DlpContentRestriction::kScreenshot);
 }  // namespace
 
-class MockDlpContentManager : public DlpContentManager {
- public:
-  MOCK_METHOD2(OnConfidentialityChanged,
-               void(content::WebContents*, const DlpContentRestrictionSet&));
-  MOCK_METHOD1(OnWebContentsDestroyed, void(const content::WebContents*));
-  MOCK_CONST_METHOD1(GetRestrictionSetForURL,
-                     DlpContentRestrictionSet(const GURL&));
-  MOCK_METHOD1(OnVisibilityChanged, void(content::WebContents*));
-};
-
 class DlpContentTabHelperTest : public ChromeRenderViewHostTestHarness {
  protected:
   void SetUp() override {

chrome/browser/chromeos/policy/dlp/mock_dlp_content_manager.cc

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chrome/browser/chromeos/policy/dlp/mock_dlp_content_manager.h"
+
+namespace policy {
+
+MockDlpContentManager::MockDlpContentManager() = default;
+
+MockDlpContentManager::~MockDlpContentManager() = default;
+
+}  // namespace policy

chrome/browser/chromeos/policy/dlp/mock_dlp_content_manager.h

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROME_BROWSER_CHROMEOS_POLICY_DLP_MOCK_DLP_CONTENT_MANAGER_H_
+#define CHROME_BROWSER_CHROMEOS_POLICY_DLP_MOCK_DLP_CONTENT_MANAGER_H_
+
+#include "chrome/browser/chromeos/policy/dlp/dlp_content_manager.h"
+#include "chrome/browser/ui/ash/chrome_screenshot_grabber.h"
+#include "testing/gmock/include/gmock/gmock.h"
+#include "url/gurl.h"
+
+namespace policy {
+
+class MockDlpContentManager : public DlpContentManager {
+ public:
+  MockDlpContentManager();
+  ~MockDlpContentManager() override;
+
+  MOCK_METHOD2(OnConfidentialityChanged,
+               void(content::WebContents*, const DlpContentRestrictionSet&));
+  MOCK_METHOD1(OnWebContentsDestroyed, void(content::WebContents*));
+  MOCK_CONST_METHOD1(GetRestrictionSetForURL,
+                     DlpContentRestrictionSet(const GURL&));
+  MOCK_METHOD1(OnVisibilityChanged, void(content::WebContents*));
+  MOCK_CONST_METHOD1(IsScreenshotRestricted, bool(const ScreenshotArea& area));
+};
+
+}  // namespace policy
+
+#endif  // CHROME_BROWSER_CHROMEOS_POLICY_DLP_MOCK_DLP_CONTENT_MANAGER_H_

chrome/browser/ui/ash/chrome_screenshot_grabber.cc

@@ -30,6 +30,7 @@
 #include "chrome/browser/browser_process.h"
 #include "chrome/browser/chromeos/file_manager/open_util.h"
 #include "chrome/browser/chromeos/note_taking_helper.h"
+#include "chrome/browser/chromeos/policy/dlp/dlp_content_manager.h"
 #include "chrome/browser/download/download_prefs.h"
 #include "chrome/browser/notifications/notification_display_service.h"
 #include "chrome/browser/platform_util.h"
@@ -272,6 +273,31 @@ void EnsureLocalDirectoryExists(
 
 }  // namespace
 
+// static
+ScreenshotArea ScreenshotArea::CreateForAllRootWindows() {
+  return ScreenshotArea(ScreenshotType::kAllRootWindows, nullptr,
+                        base::nullopt);
+}
+
+// static
+ScreenshotArea ScreenshotArea::CreateForWindow(const aura::Window* window) {
+  return ScreenshotArea(ScreenshotType::kWindow, window, base::nullopt);
+}
+
+// static
+ScreenshotArea ScreenshotArea::CreateForPartialWindow(
+    const aura::Window* window,
+    const gfx::Rect rect) {
+  return ScreenshotArea(ScreenshotType::kPartialWindow, window, rect);
+}
+
+ScreenshotArea::ScreenshotArea(const ScreenshotArea& area) = default;
+
+ScreenshotArea::ScreenshotArea(ScreenshotType type,
+                               const aura::Window* window,
+                               base::Optional<const gfx::Rect> rect)
+    : type(type), window(window), rect(rect) {}
+
 ChromeScreenshotGrabber::ChromeScreenshotGrabber()
     : screenshot_grabber_(new ui::ScreenshotGrabber) {
   DCHECK(!g_chrome_screenshot_grabber_instance);
@@ -289,7 +315,8 @@ ChromeScreenshotGrabber* ChromeScreenshotGrabber::Get() {
 }
 
 void ChromeScreenshotGrabber::HandleTakeScreenshotForAllRootWindows() {
-  if (!ScreenshotsAllowed()) {
+  const ScreenshotArea area = ScreenshotArea::CreateForAllRootWindows();
+  if (!IsScreenshotAllowed(area)) {
     OnScreenshotCompleted(ScreenshotResult::DISABLED, base::FilePath());
     return;
   }
@@ -313,7 +340,7 @@ void ChromeScreenshotGrabber::HandleTakeScreenshotForAllRootWindows() {
     screenshot_grabber_->TakeScreenshot(
         root_window, rect,
         base::BindOnce(&ChromeScreenshotGrabber::OnTookScreenshot,
-                       weak_factory_.GetWeakPtr(), time, display_id));
+                       weak_factory_.GetWeakPtr(), time, display_id, area));
   }
   base::RecordAction(base::UserMetricsAction("Screenshot_TakeFull"));
 }
@@ -321,7 +348,9 @@ void ChromeScreenshotGrabber::HandleTakeScreenshotForAllRootWindows() {
 void ChromeScreenshotGrabber::HandleTakePartialScreenshot(
     aura::Window* window,
     const gfx::Rect& rect) {
-  if (!ScreenshotsAllowed()) {
+  const ScreenshotArea area =
+      ScreenshotArea::CreateForPartialWindow(window, rect);
+  if (!IsScreenshotAllowed(area)) {
     OnScreenshotCompleted(ScreenshotResult::DISABLED, base::FilePath());
     return;
   }
@@ -330,12 +359,13 @@ void ChromeScreenshotGrabber::HandleTakePartialScreenshot(
       window, rect,
       base::BindOnce(&ChromeScreenshotGrabber::OnTookScreenshot,
                      weak_factory_.GetWeakPtr(), base::Time::Now(),
-                     base::Optional<int>()));
+                     base::Optional<int>(), area));
   base::RecordAction(base::UserMetricsAction("Screenshot_TakePartial"));
 }
 
 void ChromeScreenshotGrabber::HandleTakeWindowScreenshot(aura::Window* window) {
-  if (!ScreenshotsAllowed()) {
+  const ScreenshotArea area = ScreenshotArea::CreateForWindow(window);
+  if (!IsScreenshotAllowed(area)) {
     OnScreenshotCompleted(ScreenshotResult::DISABLED, base::FilePath());
     return;
   }
@@ -344,7 +374,7 @@ void ChromeScreenshotGrabber::HandleTakeWindowScreenshot(aura::Window* window) {
       window, gfx::Rect(window->bounds().size()),
       base::BindOnce(&ChromeScreenshotGrabber::OnTookScreenshot,
                      weak_factory_.GetWeakPtr(), base::Time::Now(),
-                     base::Optional<int>()));
+                     base::Optional<int>(), area));
   base::RecordAction(base::UserMetricsAction("Screenshot_TakeWindow"));
 }
 
@@ -355,6 +385,7 @@ bool ChromeScreenshotGrabber::CanTakeScreenshot() {
 void ChromeScreenshotGrabber::OnTookScreenshot(
     const base::Time& screenshot_time,
     const base::Optional<int>& display_num,
+    const ScreenshotArea& area,
     ScreenshotResult result,
     scoped_refptr<base::RefCountedMemory> png_data) {
   if (result != ScreenshotResult::SUCCESS) {
@@ -363,7 +394,7 @@ void ChromeScreenshotGrabber::OnTookScreenshot(
     return;
   }
 
-  if (!ScreenshotsAllowed()) {
+  if (!IsScreenshotAllowed(area)) {
     OnScreenshotCompleted(ScreenshotResult::DISABLED, base::FilePath());
     return;
   }
@@ -534,13 +565,17 @@ Profile* ChromeScreenshotGrabber::GetProfile() {
   return ProfileManager::GetActiveUserProfile();
 }
 
-bool ChromeScreenshotGrabber::ScreenshotsAllowed() const {
-  // Have two ways to disable screenshots:
+bool ChromeScreenshotGrabber::IsScreenshotAllowed(
+    const ScreenshotArea& area) const {
+  // Have three ways to disable screenshots:
   // - local state pref whose value is set from policy;
   // - simple flag which is set/unset when entering/exiting special modes where
   // screenshots should be disabled (pref is problematic because it's kept
   // across reboots, hence if the device crashes it may get stuck with the wrong
   // value).
-  return screenshots_allowed_ && !g_browser_process->local_state()->GetBoolean(
-      prefs::kDisableScreenshots);
+  // - because of DLP restricted content present in the area of the screenshot.
+  return screenshots_allowed_ &&
+         !g_browser_process->local_state()->GetBoolean(
+             prefs::kDisableScreenshots) &&
+         !policy::DlpContentManager::Get()->IsScreenshotRestricted(area);
 }

chrome/browser/ui/ash/chrome_screenshot_grabber.h

@@ -12,6 +12,7 @@
 #include "base/macros.h"
 #include "base/memory/weak_ptr.h"
 #include "base/observer_list.h"
+#include "base/optional.h"
 #include "build/build_config.h"
 #include "ui/gfx/image/image.h"
 #include "ui/message_center/public/cpp/notification.h"
@@ -38,6 +39,34 @@ enum class ScreenshotFileResult {
   CREATE_FAILED
 };
 
+// Type of the screenshot mode.
+enum class ScreenshotType {
+  kAllRootWindows,
+  kPartialWindow,
+  kWindow,
+};
+
+// Structure representing the area of screenshot.
+// For kWindow screenshots |window| should be set.
+// For kPartialWindow screenshots |rect| and |window| should be set.
+struct ScreenshotArea {
+  static ScreenshotArea CreateForAllRootWindows();
+  static ScreenshotArea CreateForWindow(const aura::Window* window);
+  static ScreenshotArea CreateForPartialWindow(const aura::Window* window,
+                                               const gfx::Rect rect);
+
+  ScreenshotArea(const ScreenshotArea& area);
+
+  const ScreenshotType type;
+  const aura::Window* window = nullptr;
+  const base::Optional<const gfx::Rect> rect;
+
+ private:
+  ScreenshotArea(ScreenshotType type,
+                 const aura::Window* window,
+                 base::Optional<const gfx::Rect> rect);
+};
+
 class ChromeScreenshotGrabber : public ash::ScreenshotDelegate {
  public:
   // Callback called with the |result| of trying to create a local writable
@@ -66,6 +95,7 @@ class ChromeScreenshotGrabber : public ash::ScreenshotDelegate {
   // callback from ScreenshotGrabber.
   void OnTookScreenshot(const base::Time& screenshot_time,
                         const base::Optional<int>& display_num,
+                        const ScreenshotArea& area,
                         ui::ScreenshotResult result,
                         scoped_refptr<base::RefCountedMemory> png_data);
 
@@ -119,7 +149,11 @@ class ChromeScreenshotGrabber : public ash::ScreenshotDelegate {
 
   Profile* GetProfile();
 
-  bool ScreenshotsAllowed() const;
+  bool IsScreenshotAllowed(const ScreenshotArea& area) const;
+
+  // Checks whether screenshots are restricted due to current DLP policy.
+  // |window| might be nullptr for full- or partial-screenshots.
+  bool IsScreenshotRestricted(aura::Window* window) const;
 
   std::unique_ptr<ui::ScreenshotGrabber> screenshot_grabber_;
 

chrome/browser/ui/ash/chrome_screenshot_grabber_browsertest.cc

@@ -6,6 +6,7 @@
 #include "base/bind.h"
 #include "base/threading/thread_restrictions.h"
 #include "chrome/browser/browser_process.h"
+#include "chrome/browser/chromeos/policy/dlp/mock_dlp_content_manager.h"
 #include "chrome/browser/chromeos/profiles/profile_helper.h"
 #include "chrome/browser/notifications/notification_display_service_tester.h"
 #include "chrome/browser/profiles/profile_manager.h"
@@ -24,6 +25,9 @@
 #include "ui/message_center/public/cpp/notification_types.h"
 #include "ui/message_center/public/cpp/notifier_id.h"
 
+using testing::_;
+using testing::Return;
+
 class ChromeScreenshotGrabberBrowserTest
     : public InProcessBrowserTest,
       public ChromeScreenshotGrabberTestObserver,
@@ -38,6 +42,8 @@ class ChromeScreenshotGrabberBrowserTest
     display_service_->SetNotificationAddedClosure(base::BindRepeating(
         &ChromeScreenshotGrabberBrowserTest::OnNotificationAdded,
         base::Unretained(this)));
+    policy::DlpContentManager::SetDlpContentManagerForTesting(
+        &mock_dlp_content_manager_);
   }
 
   void SetTestObserver(ChromeScreenshotGrabber* chrome_screenshot_grabber,
@@ -82,6 +88,7 @@ class ChromeScreenshotGrabberBrowserTest
   base::FilePath screenshot_path_;
   bool notification_added_ = false;
   bool clipboard_changed_ = false;
+  policy::MockDlpContentManager mock_dlp_content_manager_;
 
  private:
   DISALLOW_COPY_AND_ASSIGN(ChromeScreenshotGrabberBrowserTest);
@@ -146,3 +153,23 @@ IN_PROC_BROWSER_TEST_F(ChromeScreenshotGrabberBrowserTest,
   EXPECT_TRUE(display_service_->GetNotification(std::string("screenshot")));
   EXPECT_EQ(ui::ScreenshotResult::DISABLED, screenshot_result_);
 }
+
+IN_PROC_BROWSER_TEST_F(ChromeScreenshotGrabberBrowserTest,
+                       ScreenshotsRestricted) {
+  ChromeScreenshotGrabber* chrome_screenshot_grabber =
+      ChromeScreenshotGrabber::Get();
+  SetTestObserver(chrome_screenshot_grabber, this);
+
+  auto* window = ash::Shell::GetPrimaryRootWindow();
+
+  EXPECT_CALL(mock_dlp_content_manager_, IsScreenshotRestricted(_))
+      .Times(1)
+      .WillOnce(Return(true));
+
+  chrome_screenshot_grabber->HandleTakeWindowScreenshot(window);
+  RunLoop();
+
+  EXPECT_TRUE(notification_added_);
+  EXPECT_TRUE(display_service_->GetNotification(std::string("screenshot")));
+  EXPECT_EQ(ui::ScreenshotResult::DISABLED, screenshot_result_);
+}

chrome/test/BUILD.gn

@@ -2438,6 +2438,9 @@ if (!is_android) {
         "../browser/chromeos/policy/device_system_use_24hour_clock_browsertest.cc",
         "../browser/chromeos/policy/display_resolution_handler_browsertest.cc",
         "../browser/chromeos/policy/display_rotation_default_handler_browsertest.cc",
+        "../browser/chromeos/policy/dlp/dlp_content_manager_browsertest.cc",
+        "../browser/chromeos/policy/dlp/mock_dlp_content_manager.cc",
+        "../browser/chromeos/policy/dlp/mock_dlp_content_manager.h",
         "../browser/chromeos/policy/external_data_handlers/device_wilco_dtc_configuration_external_data_handler_browsertest.cc",
         "../browser/chromeos/policy/force_maximize_on_first_run_chromeos_browsertest.cc",
         "../browser/chromeos/policy/login_policy_test_base.cc",