chromeos: gets MultiUserWindowManager working with mus windows

This routes changes to MultiUserWindowManager that impact mus windows over
a mojom. This way the ash side can map these windows appropriately.

BUG=756085,875111
TEST=covered by test

Change-Id: Ie69732f290d919e591505234669328e69cc4afc9
Reviewed-on: https://chromium-review.googlesource.com/c/1330719
Commit-Queue: Scott Violet <sky@chromium.org>
Reviewed-by: Tom Sepez <tsepez@chromium.org>
Reviewed-by: Ahmed Fakhry <afakhry@chromium.org>
Cr-Commit-Position: refs/heads/master@{#609338}

ash/BUILD.gn

@@ -1264,6 +1264,8 @@ component("ash") {
     "ws/ash_gpu_interface_provider.h",
     "ws/ash_window_manager.cc",
     "ws/ash_window_manager.h",
+    "ws/multi_user_window_manager_bridge.cc",
+    "ws/multi_user_window_manager_bridge.h",
     "ws/window_lookup.cc",
     "ws/window_service_delegate_impl.cc",
     "ws/window_service_delegate_impl.h",

ash/multi_user/multi_user_window_manager.cc

@@ -15,12 +15,13 @@
 #include "ash/session/session_controller.h"
 #include "ash/shell.h"
 #include "ash/wm/tablet_mode/tablet_mode_controller.h"
+#include "ash/ws/window_service_owner.h"
 #include "base/auto_reset.h"
 #include "base/macros.h"
-#include "ui/aura/client/aura_constants.h"
+#include "services/ws/window_service.h"
+#include "ui/aura/mus/window_mus.h"
+#include "ui/aura/mus/window_tree_client.h"
 #include "ui/aura/window.h"
-#include "ui/aura/window_event_dispatcher.h"
-#include "ui/aura/window_tree_host.h"
 #include "ui/base/ui_base_types.h"
 #include "ui/events/event.h"
 #include "ui/views/mus/mus_client.h"
@@ -82,6 +83,35 @@ mojom::WallpaperUserInfoPtr WallpaperUserInfoForAccount(
   return wallpaper_user_info;
 }
 
+// If |window| has a remote client, this converts it to the remote window used
+// by the delegate. This effectively undoes the mapping that
+// MultiUserWindowManagerBridge does.
+// TODO: remove this and instead notify about changes to these windows over a
+// mojom. https://crbug.com/875111.
+aura::Window* MapWindowIfNecessary(aura::Window* window) {
+  if (!ws::WindowService::HasRemoteClient(window) ||
+      !views::MusClient::Exists()) {
+    return window;
+  }
+
+  const ws::Id window_id = Shell::Get()
+                               ->window_service_owner()
+                               ->window_service()
+                               ->GetTopLevelWindowId(window);
+  if (window_id == ws::kInvalidTransportId)
+    return window;
+
+  // children[0] is used to deal with DesktopNativeWidgetAura. In particular,
+  // client code generally expects to see the first child, which corresponds to
+  // Widget::GetNativeWindow(), when using DesktopNativeWidgetAura.
+  aura::WindowMus* window_mus =
+      views::MusClient::Get()->window_tree_client()->GetWindowByServerId(
+          window_id);
+  return window_mus && !window_mus->GetWindow()->children().empty()
+             ? window_mus->GetWindow()->children()[0]
+             : window;
+}
+
 }  // namespace
 
 // A class to temporarily change the animation properties for a window.
@@ -154,7 +184,8 @@ MultiUserWindowManager* MultiUserWindowManager::Get() {
 }
 
 void MultiUserWindowManager::SetWindowOwner(aura::Window* window,
-                                            const AccountId& account_id) {
+                                            const AccountId& account_id,
+                                            bool show_for_current_user) {
   // Make sure the window is valid and there was no owner yet.
   DCHECK(window);
   DCHECK(account_id.is_valid());
@@ -173,7 +204,7 @@ void MultiUserWindowManager::SetWindowOwner(aura::Window* window,
 
   // Check if this window was created due to a user interaction. If it was,
   // transfer it to the current user.
-  if (window->GetProperty(aura::client::kCreatedByUserGesture))
+  if (show_for_current_user)
     window_to_entry_[window]->set_show_for_user(current_account_id_);
 
   // Add all transient children to our set of windows. Note that the function
@@ -387,8 +418,10 @@ bool MultiUserWindowManager::ShowWindowForUserIntern(
   }
 
   // Notify entry change.
-  if (delegate_)
-    delegate_->OnOwnerEntryChanged(window, account_id, minimized, teleported);
+  if (delegate_) {
+    delegate_->OnOwnerEntryChanged(MapWindowIfNecessary(window), account_id,
+                                   minimized, teleported);
+  }
   return true;
 }
 

ash/multi_user/multi_user_window_manager.h

@@ -64,8 +64,11 @@ class ASH_EXPORT MultiUserWindowManager : public SessionObserver,
 
   // Associates a window with a particular account. This may result in hiding
   // |window|. This should *not* be called more than once with a different
-  // account.
-  void SetWindowOwner(aura::Window* window, const AccountId& account_id);
+  // account. If |show_for_current_user| is true, this sets the 'shown'
+  // account to the current account.
+  void SetWindowOwner(aura::Window* window,
+                      const AccountId& account_id,
+                      bool show_for_current_user);
 
   // Sets the 'shown' account for a window. See class description for details on
   // what the 'shown' account is. This function may trigger changing the active

ash/public/interfaces/BUILD.gn

@@ -42,6 +42,7 @@ mojom("interfaces_internal") {
     "login_user_info.mojom",
     "media.mojom",
     "menu.mojom",
+    "multi_user_window_manager.mojom",
     "new_window.mojom",
     "night_light_controller.mojom",
     "note_taking_controller.mojom",

ash/public/interfaces/multi_user_window_manager.mojom

+// Copyright 2018 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.
+
+module ash.mojom;
+
+import "components/account_id/interfaces/account_id.mojom";
+
+// Used to assign windows to user accounts so that ash shows the appropriate set
+// of windows based on the active user.
+interface MultiUserWindowManager {
+  // Associates a window with an account. If |show_for_current_user| is true,
+  // the window is associated with |account_id|, but is shown for the currently
+  // active user.
+  SetWindowOwner(uint64 window_id,
+                 signin.mojom.AccountId account_id,
+                 bool show_for_current_user);
+
+  // Shows a previously registered window for the specified account.
+  ShowWindowForUser(uint64 window_id,
+                    signin.mojom.AccountId account_id);
+};

ash/ws/multi_user_window_manager_bridge.cc

+// Copyright 2018 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 "ash/ws/multi_user_window_manager_bridge.h"
+
+#include "ash/multi_user/multi_user_window_manager.h"
+#include "services/ws/window_tree.h"
+#include "ui/aura/window.h"
+
+namespace ash {
+
+MultiUserWindowManagerBridge::MultiUserWindowManagerBridge(
+    ws::WindowTree* window_tree,
+    mojo::ScopedInterfaceEndpointHandle handle)
+    : window_tree_(window_tree),
+      binding_(this,
+               mojo::AssociatedInterfaceRequest<mojom::MultiUserWindowManager>(
+                   std::move(handle))) {}
+
+MultiUserWindowManagerBridge::~MultiUserWindowManagerBridge() = default;
+
+void MultiUserWindowManagerBridge::SetWindowOwner(ws::Id window_id,
+                                                  const AccountId& account_id,
+                                                  bool show_for_current_user) {
+  // At this time this is only called once MultiUserWindowManager has been
+  // created. This needs to be fixed for the multi-process case.
+  // http://crbug.com/875111.
+  DCHECK(ash::MultiUserWindowManager::Get());
+  aura::Window* window = window_tree_->GetWindowByTransportId(window_id);
+  if (window && window_tree_->IsTopLevel(window)) {
+    ash::MultiUserWindowManager::Get()->SetWindowOwner(window, account_id,
+                                                       show_for_current_user);
+  } else {
+    DVLOG(1) << "SetWindowOwner passed invalid window, id=" << window_id;
+  }
+}
+
+void MultiUserWindowManagerBridge::ShowWindowForUser(
+    ws::Id window_id,
+    const AccountId& account_id) {
+  // At this time this is only called once MultiUserWindowManager has been
+  // created. This needs to be fixed for the multi-process case.
+  // http://crbug.com/875111.
+  DCHECK(ash::MultiUserWindowManager::Get());
+  aura::Window* window = window_tree_->GetWindowByTransportId(window_id);
+  if (window && window_tree_->IsTopLevel(window))
+    ash::MultiUserWindowManager::Get()->ShowWindowForUser(window, account_id);
+  else
+    DVLOG(1) << "ShowWindowForUser passed invalid window, id=" << window_id;
+}
+
+}  // namespace ash

ash/ws/multi_user_window_manager_bridge.h

+// Copyright 2018 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 ASH_WS_MULTI_USER_WINDOW_MANAGER_BRIDGE_H_
+#define ASH_WS_MULTI_USER_WINDOW_MANAGER_BRIDGE_H_
+
+#include "ash/public/interfaces/multi_user_window_manager.mojom.h"
+#include "mojo/public/cpp/bindings/associated_binding.h"
+#include "services/ws/common/types.h"
+#include "services/ws/window_manager_interface.h"
+
+namespace mojo {
+class ScopedInterfaceEndpointHandle;
+}
+
+namespace ws {
+class WindowTree;
+}
+
+namespace ash {
+
+// Trivially forwards calls to MultiUserWindowManager.
+class MultiUserWindowManagerBridge : public mojom::MultiUserWindowManager,
+                                     public ws::WindowManagerInterface {
+ public:
+  MultiUserWindowManagerBridge(ws::WindowTree* window_tree,
+                               mojo::ScopedInterfaceEndpointHandle handle);
+  ~MultiUserWindowManagerBridge() override;
+
+  // mojom::MultiUserWindowManager overrides:
+  void SetWindowOwner(ws::Id window_id,
+                      const AccountId& account_id,
+                      bool show_for_current_user) override;
+  void ShowWindowForUser(ws::Id window_id,
+                         const AccountId& account_id) override;
+
+ private:
+  ws::WindowTree* window_tree_;
+  mojo::AssociatedBinding<mojom::MultiUserWindowManager> binding_;
+
+  DISALLOW_COPY_AND_ASSIGN(MultiUserWindowManagerBridge);
+};
+
+}  // namespace ash
+
+#endif  // ASH_WS_MULTI_USER_WINDOW_MANAGER_BRIDGE_H_

ash/ws/window_lookup.cc

@@ -18,12 +18,5 @@ aura::Window* GetWindowByClientId(ws::Id transport_id) {
       ->GetWindowByClientId(transport_id);
 }
 
-ws::ClientSpecificId GetFirstWindowTreeClientId() {
-  return Shell::Get()
-      ->window_service_owner()
-      ->window_service()
-      ->GetFirstWindowTreeClientId();
-}
-
 }  // namespace window_lookup
 }  // namespace ash

ash/ws/window_lookup.h

@@ -20,10 +20,6 @@ namespace window_lookup {
 // Returns the aura::Window by transport id.
 ASH_EXPORT aura::Window* GetWindowByClientId(ws::Id transport_id);
 
-// Returns the id of the first WindowTreeClient. That is, the id assigned to
-// the first client that connects to the WindowService.
-ASH_EXPORT ws::ClientSpecificId GetFirstWindowTreeClientId();
-
 }  // namespace window_lookup
 }  // namespace ash
 

ash/ws/window_service_delegate_impl.cc

@@ -16,6 +16,7 @@
 #include "ash/wm/window_finder.h"
 #include "ash/wm/window_util.h"
 #include "ash/ws/ash_window_manager.h"
+#include "ash/ws/multi_user_window_manager_bridge.h"
 #include "base/bind.h"
 #include "mojo/public/cpp/bindings/map.h"
 #include "services/ws/public/mojom/window_manager.mojom.h"
@@ -230,10 +231,14 @@ WindowServiceDelegateImpl::CreateWindowManagerInterface(
     ws::WindowTree* tree,
     const std::string& name,
     mojo::ScopedInterfaceEndpointHandle handle) {
-  if (name != mojom::AshWindowManager::Name_)
-    return nullptr;
+  if (name == mojom::AshWindowManager::Name_)
+    return std::make_unique<AshWindowManager>(tree, std::move(handle));
 
-  return std::make_unique<AshWindowManager>(tree, std::move(handle));
+  if (name == mojom::MultiUserWindowManager::Name_) {
+    return std::make_unique<MultiUserWindowManagerBridge>(tree,
+                                                          std::move(handle));
+  }
+  return nullptr;
 }
 
 }  // namespace ash

chrome/browser/ui/ash/multi_user/DEPS

@@ -8,5 +8,6 @@ specific_include_rules = {
   "multi_user_window_manager_chromeos_unittest\.cc": [
     "+ash/multi_user/user_switch_animator.h",
     "+ash/session/session_controller.h",
+    "+ash/ws/window_lookup.h",
   ],
 }

chrome/browser/ui/ash/multi_user/multi_user_window_manager_chromeos.cc

@@ -26,8 +26,14 @@
 #include "extensions/browser/app_window/app_window.h"
 #include "extensions/browser/app_window/app_window_registry.h"
 #include "ui/aura/client/aura_constants.h"
+#include "ui/aura/env.h"
+#include "ui/aura/mus/window_mus.h"
+#include "ui/aura/mus/window_tree_client.h"
 #include "ui/aura/window.h"
+#include "ui/base/ui_base_features.h"
 #include "ui/base/ui_base_types.h"
+#include "ui/views/mus/mus_client.h"
+#include "ui/views/widget/widget.h"
 
 namespace {
 
@@ -80,6 +86,26 @@ void RecordUMAForTransferredWindowType(aura::Window* window) {
                             NUM_TELEPORT_WINDOW_TYPES);
 }
 
+// Returns the WindowMus to use when sending messages to the server.
+aura::WindowMus* GetWindowMus(aura::Window* window) {
+  if (!aura::WindowMus::Get(window))
+    return nullptr;
+
+  aura::Window* root_window = window->GetRootWindow();
+  if (!root_window)
+    return nullptr;
+
+  // DesktopNativeWidgetAura creates two aura Windows. GetNativeWindow() returns
+  // the child window. Get the widget for |window| and its root. If the Widgets
+  // are the same, it means |window| is the native window of a
+  // DesktopNativeWidgetAura. Use the root window to notify the server as that
+  // corresponds to the top-level window that ash knows about.
+  views::Widget* widget = views::Widget::GetWidgetForNativeView(window);
+  views::Widget* root_widget =
+      views::Widget::GetWidgetForNativeView(root_window);
+  return widget == root_widget ? aura::WindowMus::Get(root_window) : nullptr;
+}
+
 }  // namespace
 
 // This class keeps track of all applications which were started for a user.
@@ -110,7 +136,14 @@ MultiUserWindowManagerChromeOS::MultiUserWindowManagerChromeOS(
     : current_account_id_(current_account_id),
       ash_multi_user_window_manager_(
           std::make_unique<ash::MultiUserWindowManager>(this,
-                                                        current_account_id)) {}
+                                                        current_account_id)) {
+  if (features::IsUsingWindowService()) {
+    multi_user_window_manager_mojom_ =
+        views::MusClient::Get()
+            ->window_tree_client()
+            ->BindWindowManagerInterface<ash::mojom::MultiUserWindowManager>();
+  }
+}
 
 MultiUserWindowManagerChromeOS::~MultiUserWindowManagerChromeOS() {
   // This may trigger callbacks to us, delete it early on.
@@ -176,14 +209,25 @@ void MultiUserWindowManagerChromeOS::SetWindowOwner(
   DCHECK(GetWindowOwner(window).empty());
   window_to_entry_[window] = new WindowEntry(account_id);
 
-  ash_multi_user_window_manager_->SetWindowOwner(window, account_id);
+  // Check if this window was created due to a user interaction. If it was,
+  // transfer it to the current user.
+  const bool show_for_current_user =
+      window->GetProperty(aura::client::kCreatedByUserGesture);
+  if (window->env()->mode() == aura::Env::Mode::MUS) {
+    aura::WindowMus* window_mus = GetWindowMus(window);
+    if (window_mus) {
+      multi_user_window_manager_mojom_->SetWindowOwner(
+          window_mus->server_id(), account_id, show_for_current_user);
+    }  // else case can happen during shutdown, or for child windows.
+  } else {
+    ash_multi_user_window_manager_->SetWindowOwner(window, account_id,
+                                                   show_for_current_user);
+  }
 
   // Add observers to track state changes.
   window->AddObserver(this);
 
-  // Check if this window was created due to a user interaction. If it was,
-  // transfer it to the current user.
-  if (window->GetProperty(aura::client::kCreatedByUserGesture))
+  if (show_for_current_user)
     window_to_entry_[window]->set_show_for_user(current_account_id_);
 
   // Notify entry adding.
@@ -200,7 +244,18 @@ const AccountId& MultiUserWindowManagerChromeOS::GetWindowOwner(
 void MultiUserWindowManagerChromeOS::ShowWindowForUser(
     aura::Window* window,
     const AccountId& account_id) {
-  ash_multi_user_window_manager_->ShowWindowForUser(window, account_id);
+  if (!window)
+    return;
+
+  if (window->env()->mode() == aura::Env::Mode::MUS) {
+    aura::WindowMus* window_mus = GetWindowMus(window);
+    if (window_mus) {
+      multi_user_window_manager_mojom_->ShowWindowForUser(
+          window_mus->server_id(), account_id);
+    }
+  } else {
+    ash_multi_user_window_manager_->ShowWindowForUser(window, account_id);
+  }
 }
 
 bool MultiUserWindowManagerChromeOS::AreWindowsSharedAmongUsers() const {

chrome/browser/ui/ash/multi_user/multi_user_window_manager_chromeos.h

@@ -10,6 +10,7 @@
 #include <string>
 
 #include "ash/multi_user/multi_user_window_manager_delegate.h"
+#include "ash/public/interfaces/multi_user_window_manager.mojom.h"
 #include "base/macros.h"
 #include "base/observer_list.h"
 #include "chrome/browser/ui/ash/multi_user/multi_user_window_manager.h"
@@ -153,8 +154,17 @@ class MultiUserWindowManagerChromeOS
   // The notification registrar to track the creation of browser windows.
   content::NotificationRegistrar registrar_;
 
+  // TODO: this won't work in the multi-process mash case. What needs to happen
+  // for the multi-process case is MultiUserWindowManagerDelegate needs to be
+  // converted to a mojom that ash uses to call this code.
+  // https://crbug.com/875111.
   std::unique_ptr<ash::MultiUserWindowManager> ash_multi_user_window_manager_;
 
+  // Only used for windows created for the window-service. For example,
+  // Browser windows when running in mash.
+  ash::mojom::MultiUserWindowManagerAssociatedPtr
+      multi_user_window_manager_mojom_;
+
   DISALLOW_COPY_AND_ASSIGN(MultiUserWindowManagerChromeOS);
 };
 

chrome/browser/ui/ash/multi_user/multi_user_window_manager_chromeos_unittest.cc

@@ -25,6 +25,7 @@
 #include "ash/wm/tablet_mode/tablet_mode_window_manager.h"
 #include "ash/wm/window_state.h"
 #include "ash/wm/wm_event.h"
+#include "ash/ws/window_lookup.h"
 #include "base/compiler_specific.h"
 #include "base/logging.h"
 #include "base/macros.h"
@@ -58,13 +59,22 @@
 #include "components/user_manager/scoped_user_manager.h"
 #include "components/user_manager/user_info.h"
 #include "components/user_manager/user_manager.h"
+#include "services/ws/common/util.h"
 #include "ui/aura/client/aura_constants.h"
+#include "ui/aura/env.h"
+#include "ui/aura/mus/window_port_mus.h"
+#include "ui/aura/mus/window_tree_client.h"
 #include "ui/aura/test/env_test_helper.h"
+#include "ui/aura/test/mus/change_completion_waiter.h"
+#include "ui/aura/test/mus/window_tree_client_test_api.h"
 #include "ui/aura/window.h"
 #include "ui/aura/window_event_dispatcher.h"
+#include "ui/base/ui_base_features.h"
 #include "ui/base/ui_base_types.h"
 #include "ui/display/manager/display_manager.h"
 #include "ui/display/test/display_manager_test_api.h"
+#include "ui/views/mus/mus_client.h"
+#include "ui/views/widget/desktop_aura/desktop_native_widget_aura.h"
 #include "ui/wm/core/window_modality_controller.h"
 #include "ui/wm/core/window_util.h"
 #include "ui/wm/public/activation_client.h"
@@ -78,6 +88,43 @@ const char kBAccountIdString[] =
 const char kArrowBAccountIdString[] =
     "->{\"account_type\":\"unknown\",\"email\":\"B\"}";
 
+aura::Window* FindWindowInAshForClientWindow(aura::Window* window) {
+  if (!aura::WindowMus::Get(window))
+    return nullptr;
+
+  // Flush all messages from the WindowTreeClient to ensure the client id has
+  // been received.
+  aura::test::WaitForAllChangesToComplete();
+
+  DCHECK_EQ(window->env()->mode(), aura::Env::Mode::MUS);
+  DCHECK(views::MusClient::Exists());
+  aura::WindowTreeClient* window_tree_client =
+      views::MusClient::Get()->window_tree_client();
+  DCHECK(window_tree_client);
+
+  aura::WindowPortMus* window_port_mus = aura::WindowPortMus::Get(window);
+  if (!window_port_mus)
+    return nullptr;  // Should only happen if we're in shutdown.
+
+  // By the time we get here the id should have been determined.
+  DCHECK(window_tree_client->id().has_value());
+  ws::ClientSpecificId client_id = *(window_tree_client->id());
+
+  // Use the top-most remote window. This should correspond to the Window of
+  // the Widget (DesktopNativeWidgetAura creates two windows).
+  const ws::Id transport_id = ws::BuildTransportId(
+      client_id,
+      static_cast<ws::ClientSpecificId>(window_port_mus->server_id()));
+  return ash::window_lookup::GetWindowByClientId(transport_id);
+}
+
+void FlushWindowTreeClientMessages() {
+  if (!views::MusClient::Exists())
+    return;
+  aura::WindowTreeClientTestApi(views::MusClient::Get()->window_tree_client())
+      .FlushForTesting();
+}
+
 const content::BrowserContext* GetActiveContext() {
   const user_manager::UserManager* user_manager =
       user_manager::UserManager::Get();
@@ -130,8 +177,14 @@ class MultiUserWindowManagerChromeOSTest : public AshTestBase {
 
  protected:
   void SwitchActiveUser(const AccountId& id) {
+    // WaitForAllChangesToComplete() does nothing in classic mode.
+    // WaitForAllChangesToComplete() is called before and after to ensure all
+    // changes have been pushed to ash before a switch, and similarly after a
+    // switch.
+    FlushWindowTreeClientMessages();
     fake_user_manager_->SwitchActiveUser(id);
     ash::MultiUserWindowManager::Get()->OnActiveUserSessionChanged(id);
+    FlushWindowTreeClientMessages();
   }
 
   // Set up the test environment for this many windows.
@@ -282,7 +335,6 @@ void MultiUserWindowManagerChromeOSTest::SetUp() {
       TestingBrowserProcess::GetGlobal()->local_state());
   ash_test_helper()->set_test_shell_delegate(new TestShellDelegateChromeOS);
   AshTestBase::SetUp();
-  // This test has mixed case user ids.
   ash_test_helper()
       ->test_session_controller_client()
       ->set_use_lower_case_user_id(false);
@@ -1533,4 +1585,159 @@ TEST_F(MultiUserWindowManagerChromeOSTest, WindowBoundsAfterTabletMode) {
   EXPECT_EQ(bounds, window(1)->bounds());
 }
 
+class MultiUserWindowManagerChromeOSMashTest
+    : public MultiUserWindowManagerChromeOSTest {
+ public:
+  MultiUserWindowManagerChromeOSMashTest() = default;
+  ~MultiUserWindowManagerChromeOSMashTest() override = default;
+
+  // MultiUserWindowManagerChromeOSTest:
+  void SetUp() override {
+    original_aura_env_mode_ =
+        aura::test::EnvTestHelper().SetMode(aura::Env::Mode::MUS);
+    feature_list_.InitWithFeatures({::features::kSingleProcessMash}, {});
+    MultiUserWindowManagerChromeOSTest::SetUp();
+    // TabletModeController calls to PowerManagerClient with a callback that is
+    // run via a posted task. Run the loop now so that we know the task is
+    // processed. Without this, the task gets processed later on, interfering
+    // with this test.
+    base::RunLoop().RunUntilIdle();
+
+    // This test configures views with mus, which means it triggers some of the
+    // DCHECKs ensuring Shell's Env is used.
+    SetRunningOutsideAsh();
+
+    // Configure views backed by mus.
+    views::MusClient::InitParams mus_client_init_params;
+    mus_client_init_params.connector =
+        ash_test_helper()->GetWindowServiceConnector();
+    mus_client_init_params.create_wm_state = false;
+    mus_client_init_params.running_in_ws_process = true;
+    mus_client_ = std::make_unique<views::MusClient>(mus_client_init_params);
+  }
+  void TearDown() override {
+    mus_client_.reset();
+    MultiUserWindowManagerChromeOSTest::TearDown();
+    aura::test::EnvTestHelper().SetMode(original_aura_env_mode_);
+  }
+
+ protected:
+  std::unique_ptr<views::Widget> CreateMusWidget() {
+    std::unique_ptr<views::Widget> widget = std::make_unique<views::Widget>();
+    views::Widget::InitParams params;
+    params.ownership = views::Widget::InitParams::WIDGET_OWNS_NATIVE_WIDGET;
+    params.bounds = gfx::Rect(0, 0, 200, 200);
+    params.native_widget =
+        mus_client_->CreateNativeWidget(params, widget.get());
+    widget->Init(params);
+    widget->Show();
+    EXPECT_EQ(aura::Env::Mode::MUS, widget->GetNativeWindow()->env()->mode());
+    // Flush all messages from the WindowTreeClient to ensure ash has finished
+    // Widget creation.
+    aura::test::WaitForAllChangesToComplete();
+    return widget;
+  }
+
+ private:
+  aura::Env::Mode original_aura_env_mode_ = aura::Env::Mode::LOCAL;
+  base::test::ScopedFeatureList feature_list_;
+  std::unique_ptr<views::MusClient> mus_client_;
+
+  DISALLOW_COPY_AND_ASSIGN(MultiUserWindowManagerChromeOSMashTest);
+};
+
+TEST_F(MultiUserWindowManagerChromeOSMashTest,
+       SingleProcessMashWindowOrdering) {
+  SetUpForThisManyWindows(1);
+
+  std::unique_ptr<views::Widget> widget = CreateMusWidget();
+  aura::Window* widget_window_in_ash =
+      FindWindowInAshForClientWindow(widget->GetNativeWindow()->parent());
+  ASSERT_TRUE(widget_window_in_ash);
+  EXPECT_EQ(widget_window_in_ash->parent(), window(0)->parent());
+
+  const AccountId account_id_A(AccountId::FromUserEmail("A"));
+  const AccountId account_id_B(AccountId::FromUserEmail("B"));
+  AddTestUser(account_id_A);
+  AddTestUser(account_id_B);
+  SwitchActiveUser(account_id_A);
+
+  // Set the windows owner.
+  ::wm::ActivationClient* activation_client =
+      ::wm::GetActivationClient(window(0)->GetRootWindow());
+  multi_user_window_manager()->SetWindowOwner(window(0), account_id_A);
+  multi_user_window_manager()->SetWindowOwner(widget->GetNativeWindow(),
+                                              account_id_A);
+  aura::test::WaitForAllChangesToComplete();
+
+  // Activate the windows one by one.
+  activation_client->ActivateWindow(widget_window_in_ash);
+  EXPECT_EQ(activation_client->GetActiveWindow(), widget_window_in_ash);
+  activation_client->ActivateWindow(window(0));
+  EXPECT_EQ(activation_client->GetActiveWindow(), window(0));
+
+  aura::Window::Windows mru_list =
+      Shell::Get()->mru_window_tracker()->BuildMruWindowList();
+  EXPECT_EQ(mru_list[0], window(0));
+  EXPECT_EQ(mru_list[1], widget_window_in_ash);
+
+  SwitchActiveUser(account_id_B);
+  EXPECT_EQ(activation_client->GetActiveWindow(), nullptr);
+
+  SwitchActiveUser(account_id_A);
+  EXPECT_EQ(activation_client->GetActiveWindow(), window(0));
+
+  mru_list = Shell::Get()->mru_window_tracker()->BuildMruWindowList();
+  ASSERT_EQ(2u, mru_list.size());
+  EXPECT_EQ(mru_list[0], window(0));
+  EXPECT_TRUE(widget_window_in_ash->Contains(mru_list[1]));
+
+  // Swap the activation order, and retry.
+  activation_client->ActivateWindow(window(0));
+  EXPECT_EQ(activation_client->GetActiveWindow(), window(0));
+  activation_client->ActivateWindow(widget_window_in_ash);
+  EXPECT_EQ(activation_client->GetActiveWindow(), widget_window_in_ash);
+
+  mru_list = Shell::Get()->mru_window_tracker()->BuildMruWindowList();
+  EXPECT_EQ(mru_list[0], widget_window_in_ash);
+  EXPECT_EQ(mru_list[1], window(0));
+
+  SwitchActiveUser(account_id_B);
+  EXPECT_EQ(activation_client->GetActiveWindow(), nullptr);
+
+  SwitchActiveUser(account_id_A);
+  EXPECT_EQ(activation_client->GetActiveWindow(), widget_window_in_ash);
+
+  mru_list = Shell::Get()->mru_window_tracker()->BuildMruWindowList();
+  EXPECT_EQ(mru_list[0], widget_window_in_ash);
+  EXPECT_EQ(mru_list[1], window(0));
+}
+
+TEST_F(MultiUserWindowManagerChromeOSMashTest, SetWindowOwner) {
+  SetUpForThisManyWindows(1);
+
+  std::unique_ptr<views::Widget> widget = CreateMusWidget();
+  const AccountId account1(AccountId::FromUserEmail("A"));
+  const AccountId account2(AccountId::FromUserEmail("B"));
+  AddTestUser(account1);
+  AddTestUser(account2);
+  SwitchActiveUser(account1);
+
+  // Make both windows be own by |account1|.
+  multi_user_window_manager()->SetWindowOwner(window(0), account1);
+  multi_user_window_manager()->SetWindowOwner(widget->GetNativeWindow(),
+                                              account1);
+  aura::test::WaitForAllChangesToComplete();
+
+  // Switch to account2, which should hide the widget.
+  SwitchActiveUser(account2);
+  EXPECT_FALSE(widget->IsVisible());
+
+  // Show the widget for the current user, which should trigger showing it.
+  multi_user_window_manager()->ShowWindowForUser(widget->GetNativeWindow(),
+                                                 account2);
+  aura::test::WaitForAllChangesToComplete();
+  EXPECT_TRUE(widget->IsVisible());
+}
+
 }  // namespace ash

services/ws/common/util.h

@@ -19,6 +19,11 @@ inline ClientSpecificId ClientWindowIdFromTransportId(Id id) {
   return static_cast<ClientSpecificId>(id & 0xFFFFFFFF);
 }
 
+inline Id BuildTransportId(ClientSpecificId connection_id,
+                           ClientSpecificId window_id) {
+  return (static_cast<Id>(connection_id) << 32) | static_cast<Id>(window_id);
+}
+
 }  // namespace ws
 
 #endif  // SERVICES_WS_COMMON_UTIL_H_

services/ws/window_service.cc

@@ -79,10 +79,6 @@ WindowService::~WindowService() {
   DCHECK(window_trees_.empty());
 }
 
-ClientSpecificId WindowService::GetFirstWindowTreeClientId() const {
-  return decrement_client_ids_ ? kInitialClientIdDecrement : kInitialClientId;
-}
-
 ServerWindow* WindowService::GetServerWindowForWindowCreateIfNecessary(
     aura::Window* window) {
   ServerWindow* server_window = ServerWindow::GetMayBeNull(window);
@@ -131,6 +127,17 @@ bool WindowService::IsTopLevelWindow(const aura::Window* window) {
   return server_window && server_window->IsTopLevel();
 }
 
+ws::Id WindowService::GetTopLevelWindowId(aura::Window* window) {
+  ServerWindow* server_window = ServerWindow::GetMayBeNull(window);
+  if (!server_window || !server_window->IsTopLevel())
+    return kInvalidTransportId;
+
+  const ws::WindowTree* owning_tree = server_window->owning_window_tree();
+  return BuildTransportId(
+      owning_tree->client_id(),
+      owning_tree->ClientWindowIdForWindow(window).sink_id());
+}
+
 aura::Window* WindowService::GetWindowByClientId(Id transport_id) {
   const ClientSpecificId client_id = ClientIdFromTransportId(transport_id);
   WindowTree* window_tree = GetTreeById(client_id);

services/ws/window_service.h

@@ -87,9 +87,6 @@ class COMPONENT_EXPORT(WINDOW_SERVICE) WindowService
                 aura::Env* env = nullptr);
   ~WindowService() override;
 
-  // Returns the id of the first WindowTreeClient.
-  ClientSpecificId GetFirstWindowTreeClientId() const;
-
   // Gets the ServerWindow for |window|, creating if necessary.
   ServerWindow* GetServerWindowForWindowCreateIfNecessary(aura::Window* window);
 
@@ -113,6 +110,10 @@ class COMPONENT_EXPORT(WINDOW_SERVICE) WindowService
   // Returns true if |window| hosts a remote client and is a toplevel window.
   static bool IsTopLevelWindow(const aura::Window* window);
 
+  // Returns the transport id for |window|. If |window| is not a top-level
+  // window, returns kInvalidTransportId.
+  ws::Id GetTopLevelWindowId(aura::Window* window);
+
   // Returns the window representing the specified id.
   aura::Window* GetWindowByClientId(Id transport_id);
 

ui/aura/mus/window_tree_client.cc

@@ -222,6 +222,11 @@ bool WindowTreeClient::WaitForDisplays() {
   return valid_wait;
 }
 
+WindowMus* WindowTreeClient::GetWindowByServerId(ws::Id id) {
+  IdToWindowMap::const_iterator it = windows_.find(id);
+  return it != windows_.end() ? it->second : nullptr;
+}
+
 void WindowTreeClient::SetCanFocus(Window* window, bool can_focus) {
   DCHECK(tree_);
   DCHECK(window);
@@ -374,11 +379,6 @@ void WindowTreeClient::RegisterWindowMus(WindowMus* window) {
   }
 }
 
-WindowMus* WindowTreeClient::GetWindowByServerId(ws::Id id) {
-  IdToWindowMap::const_iterator it = windows_.find(id);
-  return it != windows_.end() ? it->second : nullptr;
-}
-
 bool WindowTreeClient::IsWindowKnown(aura::Window* window) {
   WindowMus* window_mus = WindowMus::Get(window);
   // NOTE: this function explicitly checks for a null WindowMus as it may be

ui/aura/mus/window_tree_client.h

@@ -128,6 +128,8 @@ class AURA_EXPORT WindowTreeClient
 
   const base::Optional<uint32_t>& id() const { return id_; }
 
+  WindowMus* GetWindowByServerId(ws::Id id);
+
   void SetCanFocus(Window* window, bool can_focus);
   void SetCanAcceptDrops(WindowMus* window, bool can_accept_drops);
   void SetEventTargetingPolicy(WindowMus* window,
@@ -242,8 +244,6 @@ class AURA_EXPORT WindowTreeClient
 
   void RegisterWindowMus(WindowMus* window);
 
-  WindowMus* GetWindowByServerId(ws::Id id);
-
   bool IsWindowKnown(aura::Window* window);
 
   // Returns the oldest InFlightChange that matches |change|.

ui/aura/test/env_test_helper.h

@@ -57,6 +57,7 @@ class EnvTestHelper {
     env_->mode_ = mode;
     if (mode == Env::Mode::MUS)
       env_->EnableMusOSExchangeDataProvider();
+    env_->in_mus_shutdown_ = false;
     return old_mode;
   }
 

ui/aura/test/mus/change_completion_waiter.cc

@@ -113,7 +113,20 @@ void ChangeCompletionWaiter::OnChangeCompleted(uint32_t change_id,
 void WaitForAllChangesToComplete(WindowTreeClient* client) {
   if (Env::GetInstance()->mode() == Env::Mode::LOCAL)
     return;
-  AllChangesCompletedWaiter(client ? client : GetWindowTreeClient()).Wait();
+  if (!client)
+    client = GetWindowTreeClient();
+
+  // Wait for any local changes.
+  AllChangesCompletedWaiter(client).Wait();
+
+  // Do this to flush any incoming calls from the server.
+  WindowTreeClientTestApi(client).FlushForTesting();
+
+  // And wait for any outgoing calls.
+  AllChangesCompletedWaiter(client).Wait();
+
+  // In theory we could keep doing the above until there are no more outgoing
+  // requests.
 }
 
 }  // namespace test