1 files changed, 82 insertions, 6 deletions

diff --git a/sandbox/src/policy_target_test.cc b/sandbox/src/policy_target_test.cc
index 63948685..e4ff837 100644
--- a/sandbox/src/policy_target_test.cc
+++ b/sandbox/src/policy_target_test.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2006-2008 The Chromium Authors. All rights reserved.
+// Copyright (c) 2006-2009 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.
 
@@ -165,6 +165,12 @@ TEST(PolicyTargetTest, OpenProcess) {
 // current desktop.
 TEST(PolicyTargetTest, DesktopPolicy) {
   BrokerServices* broker = GetBroker();
+
+  // Precreate the desktop.
+  TargetPolicy* temp_policy = broker->CreatePolicy();
+  temp_policy->CreateAlternateDesktop(false);
+  temp_policy->Release();
+
   ASSERT_TRUE(broker != NULL);
 
   // Get the path to the sandboxed app.
@@ -180,7 +186,7 @@ TEST(PolicyTargetTest, DesktopPolicy) {
   PROCESS_INFORMATION target = {0};
 
   TargetPolicy* policy = broker->CreatePolicy();
-  policy->SetDesktop(L"desktop_for_sbox");
+  policy->SetAlternateDesktop(false);
   policy->SetTokenLevel(USER_INTERACTIVE, USER_LOCKDOWN);
   result = broker->SpawnTarget(prog_name, arguments.c_str(), policy, &target);
   policy->Release();
@@ -194,7 +200,8 @@ TEST(PolicyTargetTest, DesktopPolicy) {
   EXPECT_NE(::GetThreadDesktop(target.dwThreadId),
             ::GetThreadDesktop(::GetCurrentThreadId()));
 
-  HDESK desk = ::OpenDesktop(L"desktop_for_sbox", 0, FALSE, DESKTOP_ENUMERATE);
+  std::wstring desktop_name = policy->GetAlternateDesktop();
+  HDESK desk = ::OpenDesktop(desktop_name.c_str(), 0, FALSE, DESKTOP_ENUMERATE);
   EXPECT_TRUE(NULL != desk);
   EXPECT_TRUE(::CloseDesktop(desk));
   EXPECT_TRUE(::TerminateProcess(target.hProcess, 0));
@@ -204,11 +211,80 @@ TEST(PolicyTargetTest, DesktopPolicy) {
   EXPECT_TRUE(::CloseHandle(target.hProcess));
   EXPECT_TRUE(::CloseHandle(target.hThread));
 
-  // Wait for the desktop to be deleted by the destructor of TargetProcess
-  Sleep(2000);
+  // Close the desktop handle.
+  temp_policy = broker->CreatePolicy();
+  temp_policy->DestroyAlternateDesktop();
+  temp_policy->Release();
 
-  desk = ::OpenDesktop(L"desktop_for_sbox", 0, FALSE, DESKTOP_ENUMERATE);
+  // Make sure the desktop does not exist anymore.
+  desk = ::OpenDesktop(desktop_name.c_str(), 0, FALSE, DESKTOP_ENUMERATE);
   EXPECT_TRUE(NULL == desk);
 }
 
+// Launches the app in the sandbox and ask it to wait in an
+// infinite loop. Waits for 2 seconds and then check if the
+// winstation associated with the app thread is not the same as the
+// current desktop.
+TEST(PolicyTargetTest, WinstaPolicy) {
+  BrokerServices* broker = GetBroker();
+
+  // Precreate the desktop.
+  TargetPolicy* temp_policy = broker->CreatePolicy();
+  temp_policy->CreateAlternateDesktop(true);
+  temp_policy->Release();
+
+  ASSERT_TRUE(broker != NULL);
+
+  // Get the path to the sandboxed app.
+  wchar_t prog_name[MAX_PATH];
+  GetModuleFileNameW(NULL, prog_name, MAX_PATH);
+
+  std::wstring arguments(L"\"");
+  arguments += prog_name;
+  arguments += L"\" -child 0 wait";  // Don't care about the "state" argument.
+
+  // Launch the app.
+  ResultCode result = SBOX_ALL_OK;
+  PROCESS_INFORMATION target = {0};
+
+  TargetPolicy* policy = broker->CreatePolicy();
+  policy->SetAlternateDesktop(true);
+  policy->SetTokenLevel(USER_INTERACTIVE, USER_LOCKDOWN);
+  result = broker->SpawnTarget(prog_name, arguments.c_str(), policy, &target);
+  policy->Release();
+
+  EXPECT_EQ(SBOX_ALL_OK, result);
+
+  EXPECT_EQ(1, ::ResumeThread(target.hThread));
+
+  EXPECT_EQ(WAIT_TIMEOUT, ::WaitForSingleObject(target.hProcess, 2000));
+
+  EXPECT_NE(::GetThreadDesktop(target.dwThreadId),
+            ::GetThreadDesktop(::GetCurrentThreadId()));
+
+  std::wstring desktop_name = policy->GetAlternateDesktop();
+  ASSERT_FALSE(desktop_name.empty());
+
+  // Make sure there is a backslash, for the window station name.
+  EXPECT_NE(desktop_name.find_first_of(L'\\'), std::wstring::npos);
+
+  // Isolate the desktop name.
+  desktop_name = desktop_name.substr(desktop_name.find_first_of(L'\\') + 1);
+
+  HDESK desk = ::OpenDesktop(desktop_name.c_str(), 0, FALSE, DESKTOP_ENUMERATE);
+  // This should fail if the desktop is really on another window station.
+  EXPECT_FALSE(NULL != desk);
+  EXPECT_TRUE(::TerminateProcess(target.hProcess, 0));
+
+  ::WaitForSingleObject(target.hProcess, INFINITE);
+
+  EXPECT_TRUE(::CloseHandle(target.hProcess));
+  EXPECT_TRUE(::CloseHandle(target.hThread));
+
+  // Close the desktop handle.
+  temp_policy = broker->CreatePolicy();
+  temp_policy->DestroyAlternateDesktop();
+  temp_policy->Release();
+}
+
 }  // namespace sandbox