client_storage.cc 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119
  1. // Copyright 2021 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #include "components/policy/test_support/client_storage.h"
  5. #include "base/check.h"
  6. #include "base/containers/contains.h"
  7. #include "crypto/sha2.h"
  8. namespace policy {
  9. ClientStorage::ClientInfo::ClientInfo() = default;
  10. ClientStorage::ClientInfo::ClientInfo(
  11. const ClientStorage::ClientInfo& client_info) = default;
  12. ClientStorage::ClientInfo& ClientStorage::ClientInfo::operator=(
  13. const ClientStorage::ClientInfo& client_info) = default;
  14. ClientStorage::ClientInfo::ClientInfo(ClientStorage::ClientInfo&& client_info) =
  15. default;
  16. ClientStorage::ClientInfo& ClientStorage::ClientInfo::operator=(
  17. ClientStorage::ClientInfo&& client_info) = default;
  18. ClientStorage::ClientInfo::~ClientInfo() = default;
  19. ClientStorage::ClientStorage() = default;
  20. ClientStorage::ClientStorage(ClientStorage&& client_storage) = default;
  21. ClientStorage& ClientStorage::operator=(ClientStorage&& client_storage) =
  22. default;
  23. ClientStorage::~ClientStorage() = default;
  24. void ClientStorage::RegisterClient(const ClientInfo& client_info) {
  25. CHECK(!client_info.device_id.empty());
  26. clients_[client_info.device_id] = client_info;
  27. registered_tokens_[client_info.device_token] = client_info.device_id;
  28. }
  29. bool ClientStorage::HasClient(const std::string& device_id) const {
  30. return clients_.find(device_id) != clients_.end();
  31. }
  32. const ClientStorage::ClientInfo& ClientStorage::GetClient(
  33. const std::string& device_id) const {
  34. const ClientInfo* const client_info = GetClientOrNull(device_id);
  35. CHECK(client_info);
  36. return *client_info;
  37. }
  38. const ClientStorage::ClientInfo* ClientStorage::GetClientOrNull(
  39. const std::string& device_id) const {
  40. auto it = clients_.find(device_id);
  41. return it == clients_.end() ? nullptr : &it->second;
  42. }
  43. const ClientStorage::ClientInfo* ClientStorage::LookupByStateKey(
  44. const std::string& state_key) const {
  45. for (auto const& [device_id, client_info] : clients_) {
  46. if (base::Contains(client_info.state_keys, state_key))
  47. return &client_info;
  48. }
  49. return nullptr;
  50. }
  51. bool ClientStorage::DeleteClient(const std::string& device_token) {
  52. auto it = registered_tokens_.find(device_token);
  53. if (it == registered_tokens_.end())
  54. return false;
  55. const std::string& device_id = it->second;
  56. DCHECK(!device_id.empty());
  57. auto it_clients = clients_.find(device_id);
  58. DCHECK(it_clients != clients_.end());
  59. clients_.erase(it_clients, clients_.end());
  60. registered_tokens_.erase(it, registered_tokens_.end());
  61. return true;
  62. }
  63. size_t ClientStorage::GetNumberOfRegisteredClients() const {
  64. return clients_.size();
  65. }
  66. std::vector<std::string> ClientStorage::GetMatchingStateKeyHashes(
  67. uint64_t modulus,
  68. uint64_t remainder) const {
  69. std::vector<std::string> hashes;
  70. for (const auto& [device_id, client_info] : clients_) {
  71. for (const std::string& key : client_info.state_keys) {
  72. std::string hash = crypto::SHA256HashString(key);
  73. uint64_t hash_remainder = 0;
  74. // Simulate long division in base 256, which allows us to interpret
  75. // individual chars in our hash as digits. We only care about the
  76. // remainder and hence do not compute the quotient in each iteration. This
  77. // assumes big-endian byte order.
  78. for (uint64_t digit : hash)
  79. hash_remainder = (hash_remainder * 256 + digit) % modulus;
  80. if (hash_remainder == remainder)
  81. hashes.push_back(hash);
  82. }
  83. }
  84. return hashes;
  85. }
  86. std::vector<ClientStorage::ClientInfo> ClientStorage::GetAllClients() {
  87. std::vector<ClientStorage::ClientInfo> result;
  88. for (const auto& [device_id, client_info] : clients_) {
  89. result.push_back(client_info);
  90. }
  91. return result;
  92. }
  93. } // namespace policy