Refactor: Update ObjectSID and RIDSet handling to use directory IDs instead of objects for improved consistency and performance

Author: m.shvets

app/alembic/versions/552b4eafb1aa_remove_objectsid_vals.py

@@ -1,7 +1,7 @@
 """Add rIDManager and rIDSet objectClasses to LDAP schema.
 
 Revision ID: 552b4eafb1aa
-Revises: 19d86e660cf2
+Revises: df4287898910
 Create Date: 2026-02-17 09:24:57.906080
 
 """
@@ -230,10 +230,11 @@ async def _init_rid_manager(
             if rid > max_rid:
                 max_rid = rid
 
-        start_rid = max(max_rid, RIDManagerSetupUseCase.RID_USER_MIN)
+        start_rid = max(max_rid, RIDManagerSetupUseCase.RID_MIN)
 
         qword = to_qword(start_rid, RIDManagerSetupUseCase.RID_AVAILABLE_MAX)
-        await rid_setup_gateway.set_rid_available_pool(domain, qword)
+
+        await rid_setup_gateway.set_rid_available_pool(rid_manager_dir, qword)
 
         system_container = await rid_setup_gateway.get_system_container()
         await role_use_case.inherit_parent_aces(

app/alembic/versions/a1b2c3d4e5f6_rename_services_to_system.py

@@ -11,6 +11,7 @@
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
 
+from constants import SYSTEM_CONTAINER_NAME
 from entities import Attribute, Directory
 from repo.pg.tables import queryable_attr as qa
 
@@ -111,7 +112,7 @@ async def _rename_system_to_services(connection: AsyncConnection) -> None:  # no
 
         system_dir = await session.scalar(
             select(Directory).where(
-                qa(Directory.name) == "System",
+                qa(Directory.name) == SYSTEM_CONTAINER_NAME,
                 qa(Directory.is_system).is_(True),
             ),
         )

app/enums.py

@@ -293,10 +293,7 @@ class SidPrefix(StrEnum):
 class SecurityPrincipalRid(IntEnum):
     ADMINISTRATOR = 500
     GUESTS = 501
-    KRBTGT = 502
     DOMAIN_ADMINS = 512
     DOMAIN_USERS = 513
-    DOMAIN_GUESTS = 514
     DOMAIN_COMPUTERS = 515
-    DOMAIN_CONTROLLERS = 516
     DOMAIN_READ_ONLY = 521

app/extra/scripts/add_domain_controller.py

@@ -16,7 +16,7 @@
     EntityTypeUseCase,
 )
 from ldap_protocol.objects import UserAccountControlFlag
-from ldap_protocol.rid_manager import ObjectSIDUseCase
+from ldap_protocol.rid_manager import ObjectSIDUseCase, RIDSetUseCase
 from ldap_protocol.roles.role_use_case import RoleUseCase
 from repo.pg.tables import queryable_attr as qa
 
@@ -28,6 +28,7 @@ async def _add_domain_controller(
     settings: Settings,
     dc_ou_dir: Directory,
     object_sid_use_case: ObjectSIDUseCase,
+    rid_set_use_case: RIDSetUseCase,
 ) -> None:
     dc_directory = Directory(
         object_class="",
@@ -40,9 +41,13 @@ async def _add_domain_controller(
 
     dc_directory.parent_id = dc_ou_dir.id
     await object_sid_use_case.add(
-        directory=dc_directory,
+        directory_id=dc_directory.id,
     )
     await session.flush()
+    await rid_set_use_case.add(
+        domain_controller=dc_directory,
+        allocation_params=await rid_set_use_case.generate_rid_set_attrs(),
+    )
 
     attributes = [
         Attribute(
@@ -105,6 +110,7 @@ async def add_domain_controller(
     role_use_case: RoleUseCase,
     entity_type_use_case: EntityTypeUseCase,
     object_sid_use_case: ObjectSIDUseCase,
+    rid_set_use_case: RIDSetUseCase,
 ) -> None:
     logger.info("Adding domain controller.")
 
@@ -139,6 +145,7 @@ async def add_domain_controller(
         settings=settings,
         dc_ou_dir=domain_controllers_ou,
         object_sid_use_case=object_sid_use_case,
+        rid_set_use_case=rid_set_use_case,
     )
 
     logger.debug("Domain controller added.")

app/ldap_protocol/auth/setup_gateway.py

@@ -73,7 +73,6 @@ async def setup_enviroment(
     ) -> None:
         """Create directories and users for enviroment."""
         async with self._session.begin_nested():
-            self._session.add(domain)
             self._session.add(
                 NetworkPolicy(
                     name="Default open policy",
@@ -122,14 +121,6 @@ async def setup_enviroment(
             logger.error(traceback.format_exc())
             raise
 
-    async def is_base_domain_created(self) -> bool:
-        """Check if base domain is created."""
-        cat_result = await self._session.execute(select(Directory))
-        if cat_result.scalar_one_or_none():
-            logger.warning("dev data already set up")
-            return True
-        return False
-
     async def create_base_domain(
         self,
         dn: str = "multifactor.dev",
@@ -175,7 +166,7 @@ async def create_dir(
 
         if "objectSid" in data:
             await self._object_sid_use_case.add(
-                directory=dir_,
+                directory_id=dir_.id,
                 rid=int(data["objectSid"]),
                 sid_prefix=SidPrefix.BUILT_IN_DOMAIN,
             )

app/ldap_protocol/auth/use_cases.py

@@ -210,8 +210,6 @@ async def _create(self, dto: SetupDTO, data: list) -> None:
         :return: None.
         """
         try:
-            if await self._setup_gateway.is_base_domain_created():
-                return
             domain = await self._setup_gateway.create_base_domain(dto.domain)
             await self._rid_manager_setup_use_case.create_domain_identifier()
             await self._setup_gateway.setup_enviroment(
@@ -255,7 +253,7 @@ async def _create(self, dto: SetupDTO, data: list) -> None:
             await self._rid_manager_setup_use_case.setup()
             dc = await self._rid_manager_use_case.get_domain_controller()
             await self._object_sid_use_case.add(
-                directory=dc,
+                directory_id=dc.id,
             )
 
             await self._session.commit()

app/ldap_protocol/ldap_requests/add.py

@@ -104,9 +104,9 @@ def from_data(cls, data: ASN1Row) -> "AddRequest":
                 type=attr.value[0].value,
                 vals=[val.value for val in attr.value[1].value],
             )
-            for attr in attributes.value  # type: ignore
+            for attr in attributes.value
         ]
-        return cls(entry=entry.value, attributes=attributes)  # type: ignore
+        return cls(entry=entry.value, attributes=attributes)
 
     async def handle(  # noqa: C901
         self,
@@ -214,7 +214,7 @@ async def handle(  # noqa: C901
 
             await ctx.session.flush()
             await ctx.object_sid_use_case.add(
-                directory=new_dir,
+                directory_id=new_dir.id,
             )
             await ctx.session.flush()
         except IntegrityError:

app/ldap_protocol/rid_manager/object_sid_gateway.py

@@ -7,11 +7,12 @@
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from entities import Attribute, Directory
+from entities import Attribute
 from ldap_protocol.rid_manager.exceptions import (
     RIDManagerDomainIdentifierNotFoundError,
     RIDManagerObjectSIDNotFoundError,
 )
+from ldap_protocol.utils.async_cache import domain_identifier_cache
 from repo.pg.tables import queryable_attr as qa
 
 
@@ -22,11 +23,11 @@ def __init__(self, session: AsyncSession) -> None:
         """Initialize Object SID gateway."""
         self._session = session
 
-    async def get(self, directory: Directory) -> str:
+    async def get(self, directory_id: int) -> str:
         """Get object SID."""
         query = await self._session.scalar(
             select(Attribute).where(
-                qa(Attribute.directory_id) == directory.id,
+                qa(Attribute.directory_id) == directory_id,
                 qa(Attribute.name) == "objectSid",
             ),
         )
@@ -35,25 +36,26 @@ async def get(self, directory: Directory) -> str:
 
         return query.value
 
-    async def add(self, directory: Directory, object_sid: str) -> None:
+    async def add(self, directory_id: int, object_sid: str) -> None:
         """Add object SID."""
         self._session.add(
             Attribute(
                 name="objectSid",
                 value=object_sid,
-                directory_id=directory.id,
+                directory_id=directory_id,
             ),
         )
 
-    async def get_domain_identifier(self, domain: Directory) -> str:
-        """Get domain identifier.
+    async def get_domain_identifier(self) -> str:
+        """Get domain identifier (cached ``Attribute.value`` string)."""
+        return await domain_identifier_cache.get_or_load(
+            self._load_domain_identifier_value,
+        )
 
-        :return: Domain identifier
-        """
+    async def _load_domain_identifier_value(self) -> str:
         query = await self._session.scalar(
             select(Attribute).where(
                 qa(Attribute.name) == "DomainIdentifier",
-                qa(Attribute.directory_id) == domain.id,
             ),
         )
 

app/ldap_protocol/rid_manager/object_sid_use_case.py

@@ -6,12 +6,10 @@
 
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from entities import Directory
 from enums import SidPrefix
 from ldap_protocol.rid_manager.object_sid_gateway import ObjectSIDGateway
 from ldap_protocol.rid_manager.rid_manager_use_case import RIDManagerUseCase
 from ldap_protocol.rid_manager.rid_set_use_case import RIDSetUseCase
-from ldap_protocol.utils.queries import get_base_directories
 
 
 class ObjectSIDUseCase:
@@ -30,13 +28,13 @@ def __init__(
         self._session = session
         self._rid_manager_use_case = rid_manager_use_case
 
-    async def get(self, directory: Directory) -> str:
+    async def get(self, directory_id: int) -> str:
         """Get object SID."""
-        return await self._gateway.get(directory)
+        return await self._gateway.get(directory_id)
 
     async def add(
         self,
-        directory: Directory,
+        directory_id: int,
         rid: int | None = None,
         sid_prefix: SidPrefix = SidPrefix.DOMAIN_IDENTIFIER,
     ) -> None:
@@ -47,19 +45,13 @@ async def add(
             )
             rid_set = await self._rid_set_use_case.get(domain_controller)
             rid = await self._rid_set_use_case.allocate_next_rid(
-                rid_set,
+                rid_set.id,
             )
 
         if sid_prefix == SidPrefix.BUILT_IN_DOMAIN:
             object_sid = f"{sid_prefix}-{rid}"
         elif sid_prefix == SidPrefix.DOMAIN_IDENTIFIER:
-            domain_identifier = await self.get_domain_identifier()
+            domain_identifier = await self._gateway.get_domain_identifier()
             object_sid = f"{sid_prefix}-{domain_identifier}-{rid}"
 
-        await self._gateway.add(directory, object_sid)
-
-    async def get_domain_identifier(self) -> str:
-        """Get domain identifier."""
-        domain = (await get_base_directories(self._session))[0]
-
-        return await self._gateway.get_domain_identifier(domain)
+        await self._gateway.add(directory_id, object_sid)

app/ldap_protocol/rid_manager/rid_manager_gateway.py

@@ -8,6 +8,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from config import Settings
+from constants import DOMAIN_CONTROLLERS_OU_NAME
 from entities import Attribute, Directory
 from ldap_protocol.rid_manager.exceptions import (
     RIDManagerAvailablePoolNotFoundError,
@@ -59,9 +60,32 @@ async def get_domain_controller(
         self,
     ) -> Directory:
         """Get domain controller."""
+        domain = await self._session.scalar(
+            select(Directory).where(
+                qa(Directory.object_class) == "domain",
+                qa(Directory.parent_id).is_(None),
+            ),
+        )
+        if not domain:
+            raise RIDManagerDomainControllerNotFoundError(
+                "Domain controller not found",
+            )
+
+        domain_controllers_ou = await self._session.scalar(
+            select(Directory).where(
+                qa(Directory.name) == DOMAIN_CONTROLLERS_OU_NAME,
+                qa(Directory.parent_id) == domain.id,
+            ),
+        )
+        if not domain_controllers_ou:
+            raise RIDManagerDomainControllerNotFoundError(
+                "Domain controllers OU not found",
+            )
+
         domain_controller = await self._session.scalar(
             select(Directory).where(
                 qa(Directory.name) == self._settings.HOST_MACHINE_SHORT_NAME,
+                qa(Directory.parent_id) == domain_controllers_ou.id,
             ),
         )
         if not domain_controller: