cros: Make policy test server more configurable after it started

- Reloads client_state file every time
- Added possibility to save state keys for clients
- LocalPolicyTestServer provides function to store configuration file

I'm gonna introduce LocalPolicyTestServerMixin in the next CL

Bug: 934224
Change-Id: I7960159b5db8c62e2515f8cedccd469e66a4192e
Reviewed-on: https://chromium-review.googlesource.com/c/1482459
Commit-Queue: Roman Sorokin [CET] <rsorokin@chromium.org>
Reviewed-by: Pavol Marko <pmarko@chromium.org>
Reviewed-by: Achuith Bhandarkar <achuith@chromium.org>
Cr-Commit-Position: refs/heads/master@{#636061}

chrome/browser/chromeos/login/webview_login_browsertest.cc

@@ -844,7 +844,8 @@ class WebviewProxyAuthLoginTest : public WebviewLoginTest {
     // |fake_session_manager_client_| above, so the device will request policy
     // with these identifiers.
     policy_test_server_.RegisterClient(policy::PolicyBuilder::kFakeToken,
-                                       policy::PolicyBuilder::kFakeDeviceId);
+                                       policy::PolicyBuilder::kFakeDeviceId,
+                                       {} /* state_keys */);
     UpdateServedPolicyFromDevicePolicyTestHelper();
     ASSERT_TRUE(policy_test_server_.Start());
 

chrome/browser/chromeos/policy/device_cloud_policy_browsertest.cc

@@ -184,7 +184,8 @@ class KeyRotationDeviceCloudPolicyTest : public DevicePolicyCrosBrowserTest {
  private:
   void StartPolicyTestServer() {
     policy_test_server_.RegisterClient(PolicyBuilder::kFakeToken,
-                                       PolicyBuilder::kFakeDeviceId);
+                                       PolicyBuilder::kFakeDeviceId,
+                                       {} /* state_keys */);
     UpdateServedTestPolicy();
     policy_test_server_.EnableAutomaticRotationOfSigningKeys();
     EXPECT_TRUE(policy_test_server_.Start());
@@ -380,7 +381,8 @@ class SigninExtensionsDeviceCloudPolicyBrowserTest
     EXPECT_TRUE(policy_test_server_.SetSigningKeyAndSignature(
         signing_key.get(), PolicyBuilder::GetTestSigningKeySignature()));
     policy_test_server_.RegisterClient(PolicyBuilder::kFakeToken,
-                                       PolicyBuilder::kFakeDeviceId);
+                                       PolicyBuilder::kFakeDeviceId,
+                                       {} /* state_keys */);
     EXPECT_TRUE(policy_test_server_.Start());
   }
 

chrome/browser/chromeos/policy/device_local_account_browsertest.cc

@@ -463,7 +463,8 @@ class DeviceLocalAccountTest : public DevicePolicyCrosBrowserTest,
         signing_key.get(), PolicyBuilder::GetTestSigningKeySignature()));
     signing_key.reset();
     test_server_.RegisterClient(PolicyBuilder::kFakeToken,
-                                PolicyBuilder::kFakeDeviceId);
+                                PolicyBuilder::kFakeDeviceId,
+                                {} /* state_keys */);
     ASSERT_TRUE(test_server_.Start());
 
     BrowserList::AddObserver(this);

chrome/browser/chromeos/policy/user_network_configuration_updater_factory_browsertest.cc

@@ -445,7 +445,8 @@ class PolicyProvidedTrustAnchorsDeviceLocalAccountTest
         signing_key.get(), PolicyBuilder::GetTestSigningKeySignature()));
     signing_key.reset();
     policy_server_.RegisterClient(PolicyBuilder::kFakeToken,
-                                  PolicyBuilder::kFakeDeviceId);
+                                  PolicyBuilder::kFakeDeviceId,
+                                  {} /* state_keys */);
     ASSERT_TRUE(policy_server_.Start());
 
     DevicePolicyCrosBrowserTest::SetUp();

chrome/browser/policy/cloud/component_cloud_policy_browsertest.cc

@@ -116,7 +116,7 @@ class ComponentCloudPolicyTest : public extensions::ExtensionBrowserTest {
   }
 
   void SetUpInProcessBrowserTestFixture() override {
-    test_server_.RegisterClient(kDMToken, kDeviceID);
+    test_server_.RegisterClient(kDMToken, kDeviceID, {} /* state_keys */);
     EXPECT_TRUE(test_server_.UpdatePolicyData(
         dm_protocol::kChromeExtensionPolicyType, kTestExtension, kTestPolicy));
     ASSERT_TRUE(test_server_.Start());

chrome/browser/policy/cloud/machine_level_user_cloud_policy_browsertest.cc

@@ -613,7 +613,7 @@ class MachineLevelUserCloudPolicyPolicyFetchTest
     base::FilePath config_path = temp_dir_.GetPath().AppendASCII("config.json");
     base::WriteFile(config_path, kTestPolicyConfig, strlen(kTestPolicyConfig));
     test_server_ = std::make_unique<LocalPolicyTestServer>(config_path);
-    test_server_->RegisterClient(kDMToken, kClientID);
+    test_server_->RegisterClient(kDMToken, kClientID, {} /* state_keys */);
   }
 
   const std::string dm_token() const { return GetParam(); }

chrome/browser/policy/test/local_policy_test_server.cc

@@ -9,6 +9,7 @@
 
 #include <algorithm>
 #include <memory>
+#include <utility>
 #include <vector>
 
 #include "base/base_paths.h"
@@ -16,6 +17,7 @@
 #include "base/json/json_writer.h"
 #include "base/numerics/safe_conversions.h"
 #include "base/path_service.h"
+#include "base/strings/string_number_conversions.h"
 #include "base/strings/stringprintf.h"
 #include "base/threading/thread_restrictions.h"
 #include "base/values.h"
@@ -48,6 +50,7 @@ const base::FilePath::CharType kClientStateFileName[] =
 const char kClientStateKeyAllowedPolicyTypes[] = "allowed_policy_types";
 const char kClientStateKeyDeviceId[] = "device_id";
 const char kClientStateKeyDeviceToken[] = "device_token";
+const char kClientStateKeyStateKeys[] = "state_keys";
 const char kClientStateKeyMachineName[] = "machine_name";
 const char kClientStateKeyMachineId[] = "machine_id";
 
@@ -61,9 +64,7 @@ bool IsUnsafeCharacter(char c) {
 }  // namespace
 
 LocalPolicyTestServer::LocalPolicyTestServer()
-    : net::LocalTestServer(net::BaseTestServer::TYPE_HTTP, base::FilePath()) {
+    : LocalPolicyTestServer(base::FilePath()) {
-  base::ScopedAllowBlockingForTesting allow_blocking;
-  CHECK(server_data_dir_.CreateUniqueTempDir());
   config_file_ = server_data_dir_.GetPath().Append(kPolicyFileName);
 }
 
@@ -72,6 +73,7 @@ LocalPolicyTestServer::LocalPolicyTestServer(const base::FilePath& config_file)
       config_file_(config_file) {
   base::ScopedAllowBlockingForTesting allow_blocking;
   CHECK(server_data_dir_.CreateUniqueTempDir());
+  client_state_file_ = server_data_dir_.GetPath().Append(kClientStateFileName);
 }
 
 LocalPolicyTestServer::LocalPolicyTestServer(const std::string& test_name)
@@ -110,10 +112,8 @@ bool LocalPolicyTestServer::SetSigningKeyAndSignature(
   // Write the signature data.
   base::FilePath signature_file =
       server_data_dir_.GetPath().Append(kSigningKeySignatureFileName);
-  bytes_written = base::WriteFile(
+  bytes_written =
-      signature_file,
+      base::WriteFile(signature_file, signature.data(), signature.size());
-      signature.c_str(),
-      signature.size());
 
   return bytes_written == base::checked_cast<int>(signature.size());
 }
@@ -122,29 +122,45 @@ void LocalPolicyTestServer::EnableAutomaticRotationOfSigningKeys() {
   automatic_rotation_of_signing_keys_enabled_ = true;
 }
 
-void LocalPolicyTestServer::RegisterClient(const std::string& dm_token,
+bool LocalPolicyTestServer::RegisterClient(
-                                           const std::string& device_id) {
+    const std::string& dm_token,
-  CHECK(server_data_dir_.IsValid());
+    const std::string& device_id,
+    const std::vector<std::string>& state_keys) {
+  base::ScopedAllowBlockingForTesting allow_blocking;
+  CHECK(!client_state_file_.empty());
 
-  std::unique_ptr<base::DictionaryValue> client_dict(
+  base::Value client_dict(base::Value::Type::DICTIONARY);
-      new base::DictionaryValue());
+  client_dict.SetKey(kClientStateKeyDeviceId, base::Value(device_id));
-  client_dict->SetString(kClientStateKeyDeviceId, device_id);
+  client_dict.SetKey(kClientStateKeyDeviceToken, base::Value(dm_token));
-  client_dict->SetString(kClientStateKeyDeviceToken, dm_token);
+  base::Value state_keys_value(base::Value::Type::LIST);
-  client_dict->SetString(kClientStateKeyMachineName, std::string());
+  for (const auto& key : state_keys) {
-  client_dict->SetString(kClientStateKeyMachineId, std::string());
+    state_keys_value.GetList().push_back(base::Value(
+        base::ToLowerASCII(base::HexEncode(key.data(), key.size()))));
+  }
+  client_dict.SetKey(kClientStateKeyStateKeys, std::move(state_keys_value));
+  client_dict.SetKey(kClientStateKeyMachineName,
+                     base::Value(base::Value::Type::STRING));
+  client_dict.SetKey(kClientStateKeyMachineId,
+                     base::Value(base::Value::Type::STRING));
 
   // Allow all policy types for now.
-  std::unique_ptr<base::ListValue> types(new base::ListValue());
+  base::Value types(base::Value::Type::LIST);
-  types->AppendString(dm_protocol::kChromeDevicePolicyType);
+  types.GetList().emplace_back(dm_protocol::kChromeDevicePolicyType);
-  types->AppendString(dm_protocol::kChromeUserPolicyType);
+  types.GetList().emplace_back(dm_protocol::kChromeUserPolicyType);
-  types->AppendString(dm_protocol::kChromePublicAccountPolicyType);
+  types.GetList().emplace_back(dm_protocol::kChromePublicAccountPolicyType);
-  types->AppendString(dm_protocol::kChromeExtensionPolicyType);
+  types.GetList().emplace_back(dm_protocol::kChromeExtensionPolicyType);
-  types->AppendString(dm_protocol::kChromeSigninExtensionPolicyType);
+  types.GetList().emplace_back(dm_protocol::kChromeSigninExtensionPolicyType);
-  types->AppendString(dm_protocol::kChromeMachineLevelUserCloudPolicyType);
+  types.GetList().emplace_back(
-  types->AppendString(dm_protocol::kChromeMachineLevelExtensionCloudPolicyType);
+      dm_protocol::kChromeMachineLevelUserCloudPolicyType);
+  types.GetList().emplace_back(
+      dm_protocol::kChromeMachineLevelExtensionCloudPolicyType);
 
-  client_dict->Set(kClientStateKeyAllowedPolicyTypes, std::move(types));
+  client_dict.SetKey(kClientStateKeyAllowedPolicyTypes, std::move(types));
-  clients_.Set(dm_token, std::move(client_dict));
+  clients_.SetKey(dm_token, std::move(client_dict));
+  std::string json;
+  base::JSONWriter::Write(clients_, &json);
+  return base::WriteFile(client_state_file_, json.data(), json.size()) ==
+         base::checked_cast<int>(json.size());
 }
 
 bool LocalPolicyTestServer::UpdatePolicy(const std::string& type,
@@ -157,7 +173,7 @@ bool LocalPolicyTestServer::UpdatePolicy(const std::string& type,
   base::FilePath policy_file = server_data_dir_.GetPath().AppendASCII(
       base::StringPrintf("policy_%s.bin", selector.c_str()));
 
-  return base::WriteFile(policy_file, policy.c_str(), policy.size()) ==
+  return base::WriteFile(policy_file, policy.data(), policy.size()) ==
          base::checked_cast<int>(policy.size());
 }
 
@@ -171,10 +187,18 @@ bool LocalPolicyTestServer::UpdatePolicyData(const std::string& type,
   base::FilePath data_file = server_data_dir_.GetPath().AppendASCII(
       base::StringPrintf("policy_%s.data", selector.c_str()));
 
-  return base::WriteFile(data_file, data.c_str(), data.size()) ==
+  return base::WriteFile(data_file, data.data(), data.size()) ==
          base::checked_cast<int>(data.size());
 }
 
+bool LocalPolicyTestServer::SetConfig(const base::Value& config) {
+  base::ScopedAllowBlockingForTesting allow_blocking;
+  std::string json;
+  base::JSONWriter::Write(config, &json);
+  return base::WriteFile(config_file_, json.data(), json.size()) ==
+         base::checked_cast<int>(json.size());
+}
+
 GURL LocalPolicyTestServer::GetServiceURL() const {
   return GetURL("device_management");
 }
@@ -249,6 +273,8 @@ bool LocalPolicyTestServer::GenerateAdditionalArguments(
   if (!net::LocalTestServer::GenerateAdditionalArguments(arguments))
     return false;
 
+  // Possible values: ERROR, WARN, INFO, DEBUG.
+  arguments->SetString("log-level", "INFO");
   arguments->SetString("config-file", config_file_.AsUTF8Unsafe());
   if (!policy_key_.empty())
     arguments->SetString("policy-key", policy_key_.AsUTF8Unsafe());
@@ -256,21 +282,11 @@ bool LocalPolicyTestServer::GenerateAdditionalArguments(
     arguments->Set("rotate-policy-keys-automatically",
                    std::make_unique<base::Value>());
   }
-  if (server_data_dir_.IsValid()) {
+  if (server_data_dir_.IsValid())
     arguments->SetString("data-dir", server_data_dir_.GetPath().AsUTF8Unsafe());
 
-    if (!clients_.empty()) {
+  if (!client_state_file_.empty())
-      std::string json;
+    arguments->SetString("client-state", client_state_file_.AsUTF8Unsafe());
-      base::JSONWriter::Write(clients_, &json);
-      base::FilePath client_state_file =
-          server_data_dir_.GetPath().Append(kClientStateFileName);
-      if (base::WriteFile(client_state_file, json.c_str(), json.size()) !=
-          base::checked_cast<int>(json.size())) {
-        return false;
-      }
-      arguments->SetString("client-state", client_state_file.AsUTF8Unsafe());
-    }
-  }
 
   return true;
 }

chrome/browser/policy/test/local_policy_test_server.h

@@ -6,6 +6,7 @@
 #define CHROME_BROWSER_POLICY_TEST_LOCAL_POLICY_TEST_SERVER_H_
 
 #include <string>
+#include <vector>
 
 #include "base/compiler_specific.h"
 #include "base/files/file_path.h"
@@ -49,12 +50,13 @@ class LocalPolicyTestServer : public net::LocalTestServer {
   // works when the server serves from a temporary directory.
   void EnableAutomaticRotationOfSigningKeys();
 
-  // Pre-configures a registered client so the server returns policy without the
+  // Register client so the server returns policy without the client having to
-  // client having to make a registration call. This must be called before
+  // make a registration call. This could be called at any time (before or after
-  // starting the server, and only works when the server serves from a temporary
+  // starting the server). This only works when the server serves from a
-  // directory.
+  // temporary directory.
-  void RegisterClient(const std::string& dm_token,
+  bool RegisterClient(const std::string& dm_token,
-                      const std::string& device_id);
+                      const std::string& device_id,
+                      const std::vector<std::string>& state_keys);
 
   // Updates policy served by the server for a given (type, entity_id) pair.
   // This only works when the server serves from a temporary directory.
@@ -81,6 +83,9 @@ class LocalPolicyTestServer : public net::LocalTestServer {
                         const std::string& entity_id,
                         const std::string& data);
 
+  // Writes |config| into the configuration file.
+  bool SetConfig(const base::Value& config);
+
   // Gets the service URL.
   GURL GetServiceURL() const;
 
@@ -96,6 +101,7 @@ class LocalPolicyTestServer : public net::LocalTestServer {
 
   base::FilePath config_file_;
   base::FilePath policy_key_;
+  base::FilePath client_state_file_;
   base::DictionaryValue clients_;
   base::ScopedTempDir server_data_dir_;
   bool automatic_rotation_of_signing_keys_enabled_ = false;

chrome/browser/policy/test/policy_testserver.py

@@ -1219,22 +1219,29 @@ class PolicyTestServer(testserver_base.BrokenPipeHandlerMixIn,
 
     Args:
       server_address: Server host and port.
+      data_dir: Directory that contains files with signature, policy, clients
+        information.
       policy_path: Names the file to read JSON-formatted policy from.
+      client_state_file: Path to file with registered clients.
       private_key_paths: List of paths to read private keys from.
       rotate_keys_automatically: Whether the keys should be rotated in a
         round-robin fashion for each policy request (by default, either the
         key specified in the config or the first key will be used for all
         requests).
+      server_base_url: The server base URL. Used for ExternalPolicyData message.
     """
     testserver_base.StoppableHTTPServer.__init__(self, server_address,
                                                  PolicyRequestHandler)
-    self._registered_tokens = {}
     self.data_dir = data_dir
     self.policy_path = policy_path
-    self.client_state_file = client_state_file
     self.rotate_keys_automatically = rotate_keys_automatically
     self.server_base_url = server_base_url
 
+    #  _registered_tokens and client_state_file kept in sync if the file is set.
+    self._registered_tokens = {}
+    self.client_state_file = client_state_file
+    self.client_state_file_timestamp = 0
+
     self.keys = []
     if private_key_paths:
       # Load specified keys from the filesystem.
@@ -1289,13 +1296,11 @@ class PolicyTestServer(testserver_base.BrokenPipeHandlerMixIn,
       pubkey = asn1der.Sequence([ algorithm, asn1der.Bitstring(rsa_pubkey) ])
       entry['public_key'] = pubkey
 
-    # Load client state.
+    try:
-    if self.client_state_file is not None:
+      self.ReadClientStateFile()
-      try:
+    except Exception as e:
-        file_contents = open(self.client_state_file).read()
+      # Could fail if file is not written yet.
-        self._registered_tokens = json.loads(file_contents, strict=False)
+      logging.info('failed to load client state %s' % e)
-      except IOError:
-        pass
 
   def GetPolicies(self):
     """Returns the policies to be used, reloaded from the backend file every
@@ -1454,6 +1459,7 @@ class PolicyTestServer(testserver_base.BrokenPipeHandlerMixIn,
       A dictionary with information about a device or user that is registered by
       dmtoken, or None if the token is not found.
     """
+    self.ReadClientStateFile()
     return self._registered_tokens.get(dmtoken, None)
 
   def LookupByStateKey(self, state_key):
@@ -1466,6 +1472,7 @@ class PolicyTestServer(testserver_base.BrokenPipeHandlerMixIn,
       A dictionary with information about a device or user or None if there is
       no matching record.
     """
+    self.ReadClientStateFile()
     for client in self._registered_tokens.values():
       if state_key.encode('hex') in client.get('state_keys', []):
         return client
@@ -1478,6 +1485,7 @@ class PolicyTestServer(testserver_base.BrokenPipeHandlerMixIn,
     Returns:
       The list of registered clients.
     """
+    self.ReadClientStateFile()
     state_keys = sum([ c.get('state_keys', [])
                        for c in self._registered_tokens.values() ], [])
     hashed_keys = map(lambda key: hashlib.sha256(key.decode('hex')).digest(),
@@ -1496,11 +1504,25 @@ class PolicyTestServer(testserver_base.BrokenPipeHandlerMixIn,
       del self._registered_tokens[dmtoken]
       self.WriteClientState()
 
+  def ReadClientStateFile(self):
+    """ Loads _registered_tokens from client_state_file."""
+    if self.client_state_file is None:
+      return
+    file_timestamp = os.stat(self.client_state_file).st_mtime
+    if file_timestamp == self.client_state_file_timestamp:
+      return
+    logging.info('load client state')
+    file_contents = open(self.client_state_file).read()
+    self._registered_tokens = json.loads(file_contents, strict=False)
+    self.client_state_file_timestamp = file_timestamp
+
   def WriteClientState(self):
     """Writes the client state back to the file."""
     if self.client_state_file is not None:
       json_data = json.dumps(self._registered_tokens)
       open(self.client_state_file, 'w').write(json_data)
+      self.client_state_file_timestamp = os.stat(
+                                         self.client_state_file).st_mtime
 
   def GetBaseFilename(self, policy_selector):
     """Returns the base filename for the given policy_selector.