[Screen Capture] Only monitor lock events for the current session

We've received reports from enterprise customers using the
browser on a virtual machine that screen sharing (via Teams,
Google Meet, etc.) will unexpectedly end. The issue was traced
back to the browser receiving a WTS_SESSION_LOCK event, it
looked as if the user had locked their screen. In response
to this event the browser ends any active screen sharing
sessions.

The issue is that those lock events were coming from other
VMs hosted on the same server. We are not interested in
the events caused by other sessions, so we should only
register for changes to the current session. The affected
customer has verified that this fix has resolved the issue.

content::VideoCaptureManager is the only class that
observes these events, so there should be no other
side effects to this change.

Bug: 1163572
Change-Id: Icfa5629eb17298f7150dedfd650b77cb34887cfa
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2598389
Commit-Queue: Austin Orion <auorion@microsoft.com>
Reviewed-by: Charlie Reis <creis@chromium.org>
Reviewed-by: Guido Urdaneta <guidou@chromium.org>
Cr-Commit-Position: refs/heads/master@{#843328}

content/browser/screenlock_monitor/screenlock_monitor_device_source.h

@@ -40,6 +40,18 @@ class CONTENT_EXPORT ScreenlockMonitorDeviceSource
   ScreenlockMonitorDeviceSource();
   ~ScreenlockMonitorDeviceSource() override;
 
+#if defined(OS_WIN)
+  // Fake session notification registration/unregistration APIs allow us to test
+  // receiving and handling messages that look as if they are sent by other
+  // sessions, without having to create a session host and a second session.
+  using WTSRegisterSessionNotificationFunction = bool (*)(HWND hwnd,
+                                                          DWORD flags);
+  using WTSUnRegisterSessionNotificationFunction = bool (*)(HWND hwnd);
+  static void SetFakeNotificationAPIsForTesting(
+      WTSRegisterSessionNotificationFunction register_function,
+      WTSUnRegisterSessionNotificationFunction unregister_function);
+#endif  // defined(OS_WIN)
+
  private:
 #if defined(OS_WIN)
   // Represents a message-only window for screenlock message handling on Win.
@@ -49,10 +61,18 @@ class CONTENT_EXPORT ScreenlockMonitorDeviceSource
     SessionMessageWindow();
     ~SessionMessageWindow();
 
+    static void SetFakeNotificationAPIsForTesting(
+        WTSRegisterSessionNotificationFunction register_function,
+        WTSUnRegisterSessionNotificationFunction unregister_function);
+
    private:
     bool OnWndProc(UINT message, WPARAM wparam, LPARAM lparam, LRESULT* result);
     void ProcessWTSSessionLockMessage(WPARAM event_id);
 
+    static WTSRegisterSessionNotificationFunction
+        register_session_notification_function_;
+    static WTSUnRegisterSessionNotificationFunction
+        unregister_session_notification_function_;
     std::unique_ptr<base::win::MessageWindow> window_;
 
     DISALLOW_COPY_AND_ASSIGN(SessionMessageWindow);

content/browser/screenlock_monitor/screenlock_monitor_device_source_win.cc

@@ -12,6 +12,49 @@
 #include "base/win/message_window.h"
 
 namespace content {
+namespace {
+bool RegisterForSessionNotifications(HWND hwnd, DWORD flag) {
+  base::OnceClosure wts_register = base::BindOnce(
+      base::IgnoreResult(&::WTSRegisterSessionNotification), hwnd, flag);
+
+  base::ThreadPool::CreateCOMSTATaskRunner({})->PostTask(
+      FROM_HERE, std::move(wts_register));
+
+  return true;
+}
+
+bool UnregisterFromSessionNotifications(HWND hwnd) {
+  return ::WTSUnRegisterSessionNotification(hwnd);
+}
+}  // namespace
+
+// Set static member function pointers with default (non-fake) APIs.
+ScreenlockMonitorDeviceSource::WTSRegisterSessionNotificationFunction
+    ScreenlockMonitorDeviceSource::SessionMessageWindow::
+        register_session_notification_function_ =
+            &RegisterForSessionNotifications;
+
+ScreenlockMonitorDeviceSource::WTSUnRegisterSessionNotificationFunction
+    ScreenlockMonitorDeviceSource::SessionMessageWindow::
+        unregister_session_notification_function_ =
+            &UnregisterFromSessionNotifications;
+
+// static
+void ScreenlockMonitorDeviceSource::SetFakeNotificationAPIsForTesting(
+    WTSRegisterSessionNotificationFunction register_function,
+    WTSUnRegisterSessionNotificationFunction unregister_function) {
+  SessionMessageWindow::SetFakeNotificationAPIsForTesting(register_function,
+                                                          unregister_function);
+}
+
+// static
+void ScreenlockMonitorDeviceSource::SessionMessageWindow::
+    SetFakeNotificationAPIsForTesting(
+        WTSRegisterSessionNotificationFunction register_function,
+        WTSUnRegisterSessionNotificationFunction unregister_function) {
+  register_session_notification_function_ = register_function;
+  unregister_session_notification_function_ = unregister_function;
+}
 
 ScreenlockMonitorDeviceSource::SessionMessageWindow::SessionMessageWindow() {
   // Create a window for receiving session change notifications.
@@ -23,21 +66,21 @@ ScreenlockMonitorDeviceSource::SessionMessageWindow::SessionMessageWindow() {
     return;
   }
 
-  base::OnceClosure wts_register =
-      base::BindOnce(base::IgnoreResult(&::WTSRegisterSessionNotification),
-                     window_->hwnd(), NOTIFY_FOR_ALL_SESSIONS);
-
-  base::ThreadPool::CreateCOMSTATaskRunner({})->PostTask(
-      FROM_HERE, std::move(wts_register));
+  // Use NOTIFY_FOR_THIS_SESSION so we only receive events from the current
+  // session, and not from other users connected to the same session host.
+  bool registered = register_session_notification_function_(
+      window_->hwnd(), NOTIFY_FOR_THIS_SESSION);
+  DCHECK(registered);
 }
 
 ScreenlockMonitorDeviceSource::SessionMessageWindow::~SessionMessageWindow() {
   // There should be no race condition between this code and the worker thread.
-  // WTSUnRegisterSessionNotification is only called from destruction as we are
-  // in shutdown, which means no other worker threads can be running.
+  // |unregister_session_notification_function_| is only called from destruction
+  // as we are in shutdown, which means no other worker threads can be running.
   if (window_) {
-    ::WTSUnRegisterSessionNotification(window_->hwnd());
-    window_.reset();
+    bool unregistered =
+        unregister_session_notification_function_(window_->hwnd());
+    DCHECK(unregistered);
   }
 }
 

content/browser/screenlock_monitor/screenlock_monitor_device_source_win_unittest.cc

+// Copyright 2021 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 "content/browser/screenlock_monitor/screenlock_monitor_device_source.h"
+
+#include "base/test/task_environment.h"
+#include "content/browser/screenlock_monitor/screenlock_monitor.h"
+#include "content/browser/screenlock_monitor/screenlock_monitor_source.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace content {
+namespace {
+
+// TODO(crbug.com/1166275): These global state variables will likely lead to
+// issues if multiple tests are run in parallel. Use caution if adding more
+// tests to this file until crbug.com/1166275 is resolved.
+static DWORD g_flag;
+static HWND g_hwnd;
+
+// These functions mock the Windows OS component that notifies subscribers
+// when a session changes state (i.e. when it locks or unlocks).
+// This is the mock for the WTSRegisterSessionNotification API.
+bool FakeRegister(HWND hwnd, DWORD flag) {
+  // Ensure only valid flags for WTSRegisterSessionNotification are used.
+  if (flag != NOTIFY_FOR_ALL_SESSIONS && flag != NOTIFY_FOR_THIS_SESSION)
+    return false;
+
+  g_flag = flag;
+  g_hwnd = hwnd;
+  return true;
+}
+
+// Mocks the WTSUnRegisterSessionNotification API.
+bool FakeUnregister(HWND hwnd) {
+  if (hwnd != g_hwnd)
+    return false;
+
+  return true;
+}
+
+// If the recipient registered with the |NOTIFY_FOR_THIS_SESSION| flag, then
+// we should only send events from this session.
+bool ShouldSendEvent(DWORD session_id) {
+  if (g_flag == NOTIFY_FOR_ALL_SESSIONS)
+    return true;
+
+  DWORD current_session_id;
+  EXPECT_TRUE(ProcessIdToSessionId(GetCurrentProcessId(), &current_session_id));
+  return session_id == current_session_id;
+}
+
+void GenerateFakeLockEvent(DWORD session_id) {
+  if (!ShouldSendEvent(session_id))
+    return;
+
+  SendMessage(g_hwnd, WM_WTSSESSION_CHANGE, WTS_SESSION_LOCK, session_id);
+}
+
+void GenerateFakeUnlockEvent(DWORD session_id) {
+  if (!ShouldSendEvent(session_id))
+    return;
+
+  SendMessage(g_hwnd, WM_WTSSESSION_CHANGE, WTS_SESSION_UNLOCK, session_id);
+}
+}  // namespace
+
+class ScreenlockMonitorTestObserver : public ScreenlockObserver {
+ public:
+  ScreenlockMonitorTestObserver() : is_screen_locked_(false) {}
+  ~ScreenlockMonitorTestObserver() override = default;
+
+  // ScreenlockObserver callbacks.
+  void OnScreenLocked() override { is_screen_locked_ = true; }
+  void OnScreenUnlocked() override { is_screen_locked_ = false; }
+
+  bool IsScreenLocked() { return is_screen_locked_; }
+
+ private:
+  bool is_screen_locked_;
+};
+
+TEST(ScreenlockMonitorDeviceSourceWinTest, FakeSessionNotifications) {
+  base::test::TaskEnvironment task_environment;
+  ScreenlockMonitorDeviceSource::SetFakeNotificationAPIsForTesting(
+      &FakeRegister, &FakeUnregister);
+  ScreenlockMonitorDeviceSource* screenlock_monitor_source =
+      new ScreenlockMonitorDeviceSource();
+  auto screenlock_monitor = std::make_unique<ScreenlockMonitor>(
+      std::unique_ptr<ScreenlockMonitorSource>(screenlock_monitor_source));
+  ScreenlockMonitorTestObserver observer;
+  screenlock_monitor->AddObserver(&observer);
+
+  DWORD session_id;
+  EXPECT_TRUE(ProcessIdToSessionId(GetCurrentProcessId(), &session_id));
+
+  // |screenlock_monitor_source_| should ignore events from other sessions.
+  GenerateFakeLockEvent(session_id + 1);
+  base::RunLoop().RunUntilIdle();
+  EXPECT_FALSE(observer.IsScreenLocked());
+
+  GenerateFakeLockEvent(session_id);
+  base::RunLoop().RunUntilIdle();
+  EXPECT_TRUE(observer.IsScreenLocked());
+
+  GenerateFakeUnlockEvent(session_id + 1);
+  base::RunLoop().RunUntilIdle();
+  EXPECT_TRUE(observer.IsScreenLocked());
+
+  GenerateFakeUnlockEvent(session_id);
+  base::RunLoop().RunUntilIdle();
+  EXPECT_FALSE(observer.IsScreenLocked());
+}
+
+// CAUTION: adding another test here will likely cause issues. See the comment
+// at the top of this file, or crbug.com/1166275
+
+}  // namespace content

content/browser/screenlock_monitor/screenlock_monitor_unittest.cc

@@ -17,7 +17,7 @@ class ScreenlockMonitorTestSource : public ScreenlockMonitorSource {
  public:
   ScreenlockMonitorTestSource() {
     DCHECK(base::CurrentThread::Get())
-        << "ScreenlocMonitorTestSource requires a MessageLoop.";
+        << "ScreenlockMonitorTestSource requires a MessageLoop.";
   }
   ~ScreenlockMonitorTestSource() override = default;
 

content/test/BUILD.gn

@@ -2387,6 +2387,7 @@ test("content_unittests") {
       "../browser/renderer_host/direct_manipulation_win_unittest.cc",
       "../browser/renderer_host/dwrite_font_lookup_table_builder_win_unittest.cc",
       "../browser/renderer_host/dwrite_font_proxy_impl_win_unittest.cc",
+      "../browser/screenlock_monitor/screenlock_monitor_device_source_win_unittest.cc",
       "../child/dwrite_font_proxy/dwrite_font_proxy_win_unittest.cc",
       "../child/dwrite_font_proxy/font_fallback_win_unittest.cc",
       "../child/font_warmup_win_unittest.cc",