ClickHouse/dbms/Access/RoleCache.cpp
Ivan 97f2a2213e
Move all folders inside /dbms one level up (#9974)
* Move some code outside dbms/src folder
* Fix paths
2020-04-02 02:51:21 +03:00

188 lines
5.8 KiB
C++

#include <Access/RoleCache.h>
#include <Access/Role.h>
#include <Access/EnabledRolesInfo.h>
#include <Access/AccessControlManager.h>
#include <boost/container/flat_map.hpp>
namespace DB
{
namespace
{
struct CollectedRoleInfo
{
RolePtr role;
bool is_current_role = false;
bool with_admin_option = false;
};
void collectRoles(boost::container::flat_map<UUID, CollectedRoleInfo> & collected_roles,
const std::function<RolePtr(const UUID &)> & get_role_function,
const UUID & role_id,
bool is_current_role,
bool with_admin_option)
{
auto it = collected_roles.find(role_id);
if (it != collected_roles.end())
{
it->second.is_current_role |= is_current_role;
it->second.with_admin_option |= with_admin_option;
return;
}
auto role = get_role_function(role_id);
collected_roles[role_id] = CollectedRoleInfo{role, is_current_role, with_admin_option};
if (!role)
return;
for (const auto & granted_role : role->granted_roles)
collectRoles(collected_roles, get_role_function, granted_role, false, false);
for (const auto & granted_role : role->granted_roles_with_admin_option)
collectRoles(collected_roles, get_role_function, granted_role, false, true);
}
std::shared_ptr<EnabledRolesInfo> collectInfoForRoles(const boost::container::flat_map<UUID, CollectedRoleInfo> & roles)
{
auto new_info = std::make_shared<EnabledRolesInfo>();
for (const auto & [role_id, collect_info] : roles)
{
const auto & role = collect_info.role;
if (!role)
continue;
if (collect_info.is_current_role)
new_info->current_roles.emplace_back(role_id);
new_info->enabled_roles.emplace_back(role_id);
if (collect_info.with_admin_option)
new_info->enabled_roles_with_admin_option.emplace_back(role_id);
new_info->names_of_roles[role_id] = role->getName();
new_info->access.merge(role->access);
new_info->access_with_grant_option.merge(role->access_with_grant_option);
new_info->settings_from_enabled_roles.merge(role->settings);
}
return new_info;
}
}
RoleCache::RoleCache(const AccessControlManager & manager_)
: manager(manager_), cache(600000 /* 10 minutes */) {}
RoleCache::~RoleCache() = default;
std::shared_ptr<const EnabledRoles> RoleCache::getEnabledRoles(
const std::vector<UUID> & roles, const std::vector<UUID> & roles_with_admin_option)
{
std::lock_guard lock{mutex};
EnabledRoles::Params params;
params.current_roles = roles;
params.current_roles_with_admin_option = roles_with_admin_option;
auto it = enabled_roles.find(params);
if (it != enabled_roles.end())
{
auto from_cache = it->second.lock();
if (from_cache)
return from_cache;
enabled_roles.erase(it);
}
auto res = std::shared_ptr<EnabledRoles>(new EnabledRoles(params));
collectRolesInfoFor(*res);
enabled_roles.emplace(std::move(params), res);
return res;
}
void RoleCache::collectRolesInfo()
{
/// `mutex` is already locked.
std::erase_if(
enabled_roles,
[&](const std::pair<EnabledRoles::Params, std::weak_ptr<EnabledRoles>> & pr)
{
auto elem = pr.second.lock();
if (!elem)
return true; // remove from the `enabled_roles` map.
collectRolesInfoFor(*elem);
return false; // keep in the `enabled_roles` map.
});
}
void RoleCache::collectRolesInfoFor(EnabledRoles & enabled)
{
/// `mutex` is already locked.
/// Collect roles in use. That includes the current roles, the roles granted to the current roles, and so on.
boost::container::flat_map<UUID, CollectedRoleInfo> collected_roles;
auto get_role_function = [this](const UUID & id) { return getRole(id); };
for (const auto & current_role : enabled.params.current_roles)
collectRoles(collected_roles, get_role_function, current_role, true, false);
for (const auto & current_role : enabled.params.current_roles_with_admin_option)
collectRoles(collected_roles, get_role_function, current_role, true, true);
/// Collect data from the collected roles.
enabled.setRolesInfo(collectInfoForRoles(collected_roles));
}
RolePtr RoleCache::getRole(const UUID & role_id)
{
/// `mutex` is already locked.
auto role_from_cache = cache.get(role_id);
if (role_from_cache)
return role_from_cache->first;
auto subscription = manager.subscribeForChanges(role_id,
[this, role_id](const UUID &, const AccessEntityPtr & entity)
{
auto changed_role = entity ? typeid_cast<RolePtr>(entity) : nullptr;
if (changed_role)
roleChanged(role_id, changed_role);
else
roleRemoved(role_id);
});
auto role = manager.tryRead<Role>(role_id);
if (role)
{
auto cache_value = Poco::SharedPtr<std::pair<RolePtr, ext::scope_guard>>(
new std::pair<RolePtr, ext::scope_guard>{role, std::move(subscription)});
cache.add(role_id, cache_value);
return role;
}
return nullptr;
}
void RoleCache::roleChanged(const UUID & role_id, const RolePtr & changed_role)
{
std::lock_guard lock{mutex};
auto role_from_cache = cache.get(role_id);
if (!role_from_cache)
return;
role_from_cache->first = changed_role;
cache.update(role_id, role_from_cache);
collectRolesInfo();
}
void RoleCache::roleRemoved(const UUID & role_id)
{
std::lock_guard lock{mutex};
cache.remove(role_id);
collectRolesInfo();
}
}