SERVER-29184 Update RoleGraph to propagate restrictions via subordinates

3 files changed, 85 insertions, 21 deletions

diff --git a/src/mongo/db/auth/role_graph.cpp b/src/mongo/db/auth/role_graph.cpp
index 24535a244fe..56ce47f2ecb 100644
--- a/src/mongo/db/auth/role_graph.cpp
+++ b/src/mongo/db/auth/role_graph.cpp
@@ -51,6 +51,7 @@ void RoleGraph::swap(RoleGraph& other) {
     swap(this->_directPrivilegesForRole, other._directPrivilegesForRole);
     swap(this->_allPrivilegesForRole, other._allPrivilegesForRole);
     swap(this->_directRestrictionsForRole, other._directRestrictionsForRole);
+    swap(this->_allRestrictionsForRole, other._allRestrictionsForRole);
     swap(this->_allRoles, other._allRoles);
 }
 
@@ -514,36 +515,40 @@ Status RoleGraph::_recomputePrivilegeDataHelper(const RoleName& startingRole,
         unordered_set<RoleName>& currentRoleIndirectRoles =
             _roleToIndirectSubordinates[currentRole];
         currentRoleIndirectRoles.clear();
-        for (std::vector<RoleName>::const_iterator it = currentRoleDirectRoles.begin();
-             it != currentRoleDirectRoles.end();
-             ++it) {
-            currentRoleIndirectRoles.insert(*it);
+        for (const auto& role : currentRoleDirectRoles) {
+            currentRoleIndirectRoles.insert(role);
+        }
+
+        // Also clear the "all restrictions" to rebuild in loop
+        auto& currentRoleAllRestrictions = _allRestrictionsForRole[currentRole];
+        currentRoleAllRestrictions.clear();
+        auto& currentRoleDirectRestrictions = _directRestrictionsForRole[currentRole];
+        if (currentRoleDirectRestrictions) {
+            currentRoleAllRestrictions.push_back(currentRoleDirectRestrictions);
         }
 
         // Recursively add children's privileges to current role's "all privileges" vector, and
         // children's roles to current roles's "indirect roles" vector.
-        for (std::vector<RoleName>::const_iterator roleIt = currentRoleDirectRoles.begin();
-             roleIt != currentRoleDirectRoles.end();
-             ++roleIt) {
+        for (const auto& childRole : currentRoleDirectRoles) {
             // At this point, we already know that the "all privilege" set for the child is
             // correct, so add those privileges to our "all privilege" set.
-            const RoleName& childRole = *roleIt;
 
-            const PrivilegeVector& childsPrivileges = _allPrivilegesForRole[childRole];
-            for (PrivilegeVector::const_iterator privIt = childsPrivileges.begin();
-                 privIt != childsPrivileges.end();
-                 ++privIt) {
-                Privilege::addPrivilegeToPrivilegeVector(&currentRoleAllPrivileges, *privIt);
+            for (const auto& priv : _allPrivilegesForRole[childRole]) {
+                Privilege::addPrivilegeToPrivilegeVector(&currentRoleAllPrivileges, priv);
             }
 
             // We also know that the "indirect roles" for the child is also correct, so we can
             // add those roles to our "indirect roles" set.
-            const unordered_set<RoleName>& childsRoles = _roleToIndirectSubordinates[childRole];
-            for (unordered_set<RoleName>::const_iterator childsRoleIt = childsRoles.begin();
-                 childsRoleIt != childsRoles.end();
-                 ++childsRoleIt) {
-                currentRoleIndirectRoles.insert(*childsRoleIt);
-            }
+            const auto& childAllRolesToIndirectSubordinates =
+                _roleToIndirectSubordinates[childRole];
+            currentRoleIndirectRoles.insert(childAllRolesToIndirectSubordinates.begin(),
+                                            childAllRolesToIndirectSubordinates.end());
+
+            // Similarly, "indirect restrictions" are ready to append
+            const auto& childAllRestrictionsForRole = _allRestrictionsForRole[childRole];
+            currentRoleAllRestrictions.insert(currentRoleAllRestrictions.end(),
+                                              childAllRestrictionsForRole.begin(),
+                                              childAllRestrictionsForRole.end());
         }
 
         visitedRoles.insert(currentRole);
diff --git a/src/mongo/db/auth/role_graph.h b/src/mongo/db/auth/role_graph.h
index 5ce151b6928..6eae09223da 100644
--- a/src/mongo/db/auth/role_graph.h
+++ b/src/mongo/db/auth/role_graph.h
@@ -138,6 +138,16 @@ public:
     }
 
     /**
+     * Returns a vector of all restriction documents that the given role contains.
+     * This includes both the restrictions set on this role directly,
+     * as well as any restrictions inherited from the role's subordinate roles.
+     */
+    const std::vector<SharedRestrictionDocument>& getAllAuthenticationRestrictions(
+        const RoleName& role) {
+        return _allRestrictionsForRole[role];
+    }
+
+    /**
      * Returns whether or not the given role exists in the role graph.  Will implicitly
      * add the role to the graph if it is a built-in role and isn't already in the graph.
      */
@@ -324,6 +334,9 @@ private:
 
     // Maps a role name to a restriction document.
     using RestrictionDocumentMap = stdx::unordered_map<RoleName, SharedRestrictionDocument>;
+    // Maps a role name to all restriction documents from self and subordinates.
+    using RestrictionDocumentsMap =
+        stdx::unordered_map<RoleName, std::vector<SharedRestrictionDocument>>;
 
     EdgeSet _roleToSubordinates;
     unordered_map<RoleName, unordered_set<RoleName>> _roleToIndirectSubordinates;
@@ -331,6 +344,7 @@ private:
     RolePrivilegeMap _directPrivilegesForRole;
     RolePrivilegeMap _allPrivilegesForRole;
     RestrictionDocumentMap _directRestrictionsForRole;
+    RestrictionDocumentsMap _allRestrictionsForRole;
     std::set<RoleName> _allRoles;
 };
 
diff --git a/src/mongo/db/auth/role_graph_test.cpp b/src/mongo/db/auth/role_graph_test.cpp
index 98c8c13c940..2e7a1ba0eb7 100644
--- a/src/mongo/db/auth/role_graph_test.cpp
+++ b/src/mongo/db/auth/role_graph_test.cpp
@@ -807,13 +807,58 @@ TEST(RoleGraphTest, AddRoleFromDocument) {
                                                                  role,
                                                                  restriction)));
                 ASSERT_OK(graph.recomputePrivilegeData());
-                ASSERT_BSONOBJ_EQ(graph.getDirectAuthenticationRestrictions(tmpRole)->toBSON(),
-                                  restriction);
+                const auto current = graph.getDirectAuthenticationRestrictions(tmpRole);
+                ASSERT_BSONOBJ_EQ(current->toBSON(), restriction);
+                const auto gaar = graph.getAllAuthenticationRestrictions(tmpRole);
+                ASSERT_TRUE(std::any_of(
+                    gaar.begin(), gaar.end(), [current](const auto& r) { return current == r; }));
                 ASSERT_OK(graph.deleteRole(tmpRole));
             }
         }
     }
 }
 
+TEST(RoleGraphTest, AddRoleFromDocumentWithRestricitonMerge) {
+    const BSONArray roleARestrictions = BSON_ARRAY(BSON("clientSource" << BSON_ARRAY("::1/128")));
+    const BSONArray roleBRestrictions =
+        BSON_ARRAY(BSON("serverAddress" << BSON_ARRAY("127.0.0.1/8")));
+
+    RoleGraph graph;
+    ASSERT_OK(graph.addRoleFromDocument(BSON("_id"
+                                             << "dbA.roleA"
+                                             << "role"
+                                             << "roleA"
+                                             << "db"
+                                             << "dbA"
+                                             << "privileges"
+                                             << BSONArray()
+                                             << "roles"
+                                             << BSONArray()
+                                             << "authenticationRestrictions"
+                                             << roleARestrictions)));
+    ASSERT_OK(graph.addRoleFromDocument(BSON("_id"
+                                             << "dbB.roleB"
+                                             << "role"
+                                             << "roleB"
+                                             << "db"
+                                             << "dbB"
+                                             << "privileges"
+                                             << BSONArray()
+                                             << "roles"
+                                             << BSON_ARRAY(BSON("role"
+                                                                << "roleA"
+                                                                << "db"
+                                                                << "dbA"))
+                                             << "authenticationRestrictions"
+                                             << roleBRestrictions)));
+    ASSERT_OK(graph.recomputePrivilegeData());
+
+    const auto A = graph.getDirectAuthenticationRestrictions(RoleName("roleA", "dbA"));
+    const auto B = graph.getDirectAuthenticationRestrictions(RoleName("roleB", "dbB"));
+    const auto gaar = graph.getAllAuthenticationRestrictions(RoleName("roleB", "dbB"));
+    ASSERT_TRUE(std::any_of(gaar.begin(), gaar.end(), [A](const auto& r) { return A == r; }));
+    ASSERT_TRUE(std::any_of(gaar.begin(), gaar.end(), [B](const auto& r) { return B == r; }));
+}
+
 }  // namespace
 }  // namespace mongo