Fix race in rbac caches

src/Access/AccessControl.cpp

@@ -75,6 +75,7 @@ class AccessControl::ContextAccessCache
             cache.remove(params);
         }
         auto res = std::shared_ptr<ContextAccess>(new ContextAccess(access_control, params));
+        res->initialize();
         cache.add(params, res);
         return res;
     }
@@ -136,10 +137,10 @@ class AccessControl::CustomSettingsPrefixes
 AccessControl::AccessControl()
     : MultipleAccessStorage("user directories"),
       context_access_cache(std::make_unique<ContextAccessCache>(*this)),
-      role_cache(std::make_unique<RoleCache>(*this)),
-      row_policy_cache(std::make_unique<RowPolicyCache>(*this)),
-      quota_cache(std::make_unique<QuotaCache>(*this)),
-      settings_profiles_cache(std::make_unique<SettingsProfilesCache>(*this)),
+      role_cache(std::make_shared<RoleCache>(*this)),
+      row_policy_cache(std::make_shared<RowPolicyCache>(*this)),
+      quota_cache(std::make_shared<QuotaCache>(*this)),
+      settings_profiles_cache(std::make_shared<SettingsProfilesCache>(*this)),
       external_authenticators(std::make_unique<ExternalAuthenticators>()),
       custom_settings_prefixes(std::make_unique<CustomSettingsPrefixes>())
 {

src/Access/AccessControl.h

@@ -161,10 +161,10 @@ class AccessControl : public MultipleAccessStorage
     class CustomSettingsPrefixes;
 
     std::unique_ptr<ContextAccessCache> context_access_cache;
-    std::unique_ptr<RoleCache> role_cache;
-    std::unique_ptr<RowPolicyCache> row_policy_cache;
-    std::unique_ptr<QuotaCache> quota_cache;
-    std::unique_ptr<SettingsProfilesCache> settings_profiles_cache;
+    std::shared_ptr<RoleCache> role_cache;
+    std::shared_ptr<RowPolicyCache> row_policy_cache;
+    std::shared_ptr<QuotaCache> quota_cache;
+    std::shared_ptr<SettingsProfilesCache> settings_profiles_cache;
     std::unique_ptr<ExternalAuthenticators> external_authenticators;
     std::unique_ptr<CustomSettingsPrefixes> custom_settings_prefixes;
 };

src/Access/ContextAccess.cpp

@@ -143,25 +143,47 @@ namespace
 
 
 ContextAccess::ContextAccess(const AccessControl & access_control_, const Params & params_)
-    : access_control(&access_control_)
-    , params(params_)
+    : access_control(&access_control_), params(params_)
+{}
+
+
+void ContextAccess::initialize()
 {
     std::lock_guard lock{mutex};
+    if (is_full_access)
+    {
+        access = std::make_shared<AccessRights>(AccessRights::getFullAccess());
+        access_with_implicit = std::make_shared<AccessRights>(addImplicitAccessRights(*access));
+        return;
+    }
+
+    auto weak_ptr = weak_from_this();
+    assert(weak_ptr);
 
     subscription_for_user_change = access_control->subscribeForChanges(
-        *params.user_id, [this](const UUID &, const AccessEntityPtr & entity)
+        *params.user_id, [weak_ptr](const UUID &, const AccessEntityPtr & entity)
     {
-        UserPtr changed_user = entity ? typeid_cast<UserPtr>(entity) : nullptr;
-        std::lock_guard lock2{mutex};
-        setUser(changed_user);
+        auto ptr = weak_ptr.lock();
+        if (!ptr)
+            return;
+        std::lock_guard lock2{ptr->mutex};
+        ptr->setUser(typeid_cast<UserPtr>(entity));
     });
 
     setUser(access_control->read<User>(*params.user_id));
 }
 
 
+ContextAccess::~ContextAccess()
+{
+    std::lock_guard lock{mutex};
+    setUser(nullptr);
+}
+
+
 void ContextAccess::setUser(const UserPtr & user_) const
 {
+    /// `mutex` is already locked.
     user = user_;
     if (!user)
     {
@@ -195,10 +217,14 @@ void ContextAccess::setUser(const UserPtr & user_) const
 
     subscription_for_roles_changes.reset();
     enabled_roles = access_control->getEnabledRoles(current_roles, current_roles_with_admin_option);
-    subscription_for_roles_changes = enabled_roles->subscribeForChanges([this](const std::shared_ptr<const EnabledRolesInfo> & roles_info_)
+    subscription_for_roles_changes = enabled_roles->subscribeForChanges(
+        [weak_ptr = weak_from_this()](const std::shared_ptr<const EnabledRolesInfo> & roles_info_)
     {
-        std::lock_guard lock{mutex};
-        setRolesInfo(roles_info_);
+        auto ptr = weak_ptr.lock();
+        if (!ptr)
+            return;
+        std::lock_guard lock{ptr->mutex};
+        ptr->setRolesInfo(roles_info_);
     });
 
     setRolesInfo(enabled_roles->getRolesInfo());
@@ -207,6 +233,7 @@ void ContextAccess::setUser(const UserPtr & user_) const
 
 void ContextAccess::setRolesInfo(const std::shared_ptr<const EnabledRolesInfo> & roles_info_) const
 {
+    /// `mutex` is already locked.
     assert(roles_info_);
     roles_info = roles_info_;
     enabled_row_policies = access_control->getEnabledRowPolicies(
@@ -221,6 +248,7 @@ void ContextAccess::setRolesInfo(const std::shared_ptr<const EnabledRolesInfo> &
 
 void ContextAccess::calculateAccessRights() const
 {
+    /// `mutex` is already locked.
     access = std::make_shared<AccessRights>(mixAccessRightsFromUserAndRoles(*user, *roles_info));
     access_with_implicit = std::make_shared<AccessRights>(addImplicitAccessRights(*access));
 
@@ -302,8 +330,7 @@ std::shared_ptr<const ContextAccess> ContextAccess::getFullAccess()
     {
         auto full_access = std::shared_ptr<ContextAccess>(new ContextAccess);
         full_access->is_full_access = true;
-        full_access->access = std::make_shared<AccessRights>(AccessRights::getFullAccess());
-        full_access->access_with_implicit = std::make_shared<AccessRights>(addImplicitAccessRights(*full_access->access));
+        full_access->initialize();
         return full_access;
     }();
     return res;

src/Access/ContextAccess.h

@@ -7,6 +7,7 @@
 #include <base/scope_guard.h>
 #include <base/shared_ptr_helper.h>
 #include <boost/container/flat_set.hpp>
+#include <memory>
 #include <mutex>
 #include <optional>
 #include <unordered_map>
@@ -63,7 +64,7 @@ struct ContextAccessParams
 };
 
 
-class ContextAccess
+class ContextAccess : public std::enable_shared_from_this<ContextAccess>
 {
 public:
     using Params = ContextAccessParams;
@@ -156,11 +157,15 @@ class ContextAccess
     /// without any limitations. This is used for the global context.
     static std::shared_ptr<const ContextAccess> getFullAccess();
 
+    ~ContextAccess();
+
 private:
     friend class AccessControl;
     ContextAccess() {}
     ContextAccess(const AccessControl & access_control_, const Params & params_);
 
+    void initialize();
+
     void setUser(const UserPtr & user_) const;
     void setRolesInfo(const std::shared_ptr<const EnabledRolesInfo> & roles_info_) const;
     void setSettingsAndConstraints() const;

src/Access/QuotaCache.cpp

@@ -163,13 +163,17 @@ QuotaCache::QuotaCache(const AccessControl & access_control_)
 {
 }
 
-QuotaCache::~QuotaCache() = default;
+QuotaCache::~QuotaCache()
+{
+    std::lock_guard lock{mutex};
+    unloadAllQuotas();
+}
 
 
 std::shared_ptr<const EnabledQuota> QuotaCache::getEnabledQuota(const UUID & user_id, const String & user_name, const boost::container::flat_set<UUID> & enabled_roles, const Poco::Net::IPAddress & client_address, const String & forwarded_address, const String & client_key)
 {
     std::lock_guard lock{mutex};
-    ensureAllQuotasRead();
+    loadAllQuotas();
 
     EnabledQuota::Params params;
     params.user_id = user_id;
@@ -194,21 +198,24 @@ std::shared_ptr<const EnabledQuota> QuotaCache::getEnabledQuota(const UUID & use
 }
 
 
-void QuotaCache::ensureAllQuotasRead()
+void QuotaCache::loadAllQuotas()
 {
     /// `mutex` is already locked.
-    if (all_quotas_read)
+    if (all_quotas_loaded)
         return;
-    all_quotas_read = true;
+    all_quotas_loaded = true;
 
     subscription = access_control.subscribeForChanges<Quota>(
-        [&](const UUID & id, const AccessEntityPtr & entity)
-        {
-            if (entity)
-                quotaAddedOrChanged(id, typeid_cast<QuotaPtr>(entity));
-            else
-                quotaRemoved(id);
-        });
+        [weak_ptr = weak_from_this()](const UUID & id, const AccessEntityPtr & entity)
+    {
+        auto ptr = weak_ptr.lock();
+        if (!ptr)
+            return;
+        if (auto quota = typeid_cast<QuotaPtr>(entity))
+            ptr->quotaAddedOrChanged(id, quota);
+        else
+            ptr->quotaRemoved(id);
+    });
 
     for (const UUID & quota_id : access_control.findAll<Quota>())
     {
@@ -219,6 +226,16 @@ void QuotaCache::ensureAllQuotasRead()
 }
 
 
+void QuotaCache::unloadAllQuotas()
+{
+    /// `mutex` is already locked.
+    subscription.reset();
+    all_quotas.clear();
+    all_quotas_loaded = false;
+    enabled_quotas.clear();
+}
+
+
 void QuotaCache::quotaAddedOrChanged(const UUID & quota_id, const std::shared_ptr<const Quota> & new_quota)
 {
     std::lock_guard lock{mutex};
@@ -250,7 +267,6 @@ void QuotaCache::quotaRemoved(const UUID & quota_id)
 void QuotaCache::chooseQuotaToConsume()
 {
     /// `mutex` is already locked.
-
     for (auto i = enabled_quotas.begin(), e = enabled_quotas.end(); i != e;)
     {
         auto elem = i->second.lock();

src/Access/QuotaCache.h

@@ -16,7 +16,7 @@ using QuotaPtr = std::shared_ptr<const Quota>;
 struct RolesOrUsersSet;
 
 /// Stores information how much amount of resources have been consumed and how much are left.
-class QuotaCache
+class QuotaCache : public std::enable_shared_from_this<QuotaCache>
 {
 public:
     QuotaCache(const AccessControl & access_control_);
@@ -52,7 +52,8 @@ class QuotaCache
         std::unordered_map<String /* quota key */, boost::shared_ptr<const Intervals>> key_to_intervals;
     };
 
-    void ensureAllQuotasRead();
+    void loadAllQuotas();
+    void unloadAllQuotas();
     void quotaAddedOrChanged(const UUID & quota_id, const std::shared_ptr<const Quota> & new_quota);
     void quotaRemoved(const UUID & quota_id);
     void chooseQuotaToConsume();
@@ -61,7 +62,7 @@ class QuotaCache
     const AccessControl & access_control;
     mutable std::mutex mutex;
     std::unordered_map<UUID /* quota id */, QuotaInfo> all_quotas;
-    bool all_quotas_read = false;
+    bool all_quotas_loaded = false;
     scope_guard subscription;
     std::map<EnabledQuota::Params, std::weak_ptr<EnabledQuota>> enabled_quotas;
 };

src/Access/RoleCache.cpp

@@ -60,7 +60,18 @@ RoleCache::RoleCache(const AccessControl & access_control_)
     : access_control(access_control_), cache(600000 /* 10 minutes */) {}
 
 
-RoleCache::~RoleCache() = default;
+RoleCache::~RoleCache()
+{
+    unloadAllRoles();
+}
+
+
+void RoleCache::unloadAllRoles()
+{
+    std::lock_guard lock{mutex};
+    cache.clear();
+    enabled_roles.clear();
+}
 
 
 std::shared_ptr<const EnabledRoles>
@@ -92,7 +103,6 @@ RoleCache::getEnabledRoles(const std::vector<UUID> & roles, const std::vector<UU
 void RoleCache::collectEnabledRoles(scope_guard & notifications)
 {
     /// `mutex` is already locked.
-
     for (auto i = enabled_roles.begin(), e = enabled_roles.end(); i != e;)
     {
         auto elem = i->second.lock();
@@ -131,19 +141,21 @@ void RoleCache::collectEnabledRoles(EnabledRoles & enabled, scope_guard & notifi
 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 = access_control.subscribeForChanges(role_id,
-                                                    [this, role_id](const UUID &, const AccessEntityPtr & entity)
+    auto subscription = access_control.subscribeForChanges(
+        role_id,
+        [weak_ptr = weak_from_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);
+        auto ptr = weak_ptr.lock();
+        if (!ptr)
+            return;
+        if (auto role = typeid_cast<RolePtr>(entity))
+            ptr->roleChanged(role_id, role);
         else
-            roleRemoved(role_id);
+            ptr->roleRemoved(role_id);
     });
 
     auto role = access_control.tryRead<Role>(role_id);

src/Access/RoleCache.h

@@ -13,7 +13,7 @@ class AccessControl;
 struct Role;
 using RolePtr = std::shared_ptr<const Role>;
 
-class RoleCache
+class RoleCache : public std::enable_shared_from_this<RoleCache>
 {
 public:
     RoleCache(const AccessControl & access_control_);
@@ -24,6 +24,7 @@ class RoleCache
         const std::vector<UUID> & current_roles_with_admin_option);
 
 private:
+    void unloadAllRoles();
     void collectEnabledRoles(scope_guard & notifications);
     void collectEnabledRoles(EnabledRoles & enabled, scope_guard & notifications);
     RolePtr getRole(const UUID & role_id);