From 9b4ff5641a583ca4fa305165e460213d8201b2d9 Mon Sep 17 00:00:00 2001
From: Helmut Grohne <helmut@subdivi.de>
Date: Mon, 27 May 2024 17:01:59 +0200
Subject: add IDAllocation.reserve method

Allow reserving a particular range instead of allocating a suitable
large range of an IDAllocation. This is useful when a directory
hierarchy defines the allocation and we merely want to verify it to be
assigned.
---
 linuxnamespaces/__init__.py | 35 +++++++++++++++++++++++++++++++++++
 tests/test_simple.py        | 18 ++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/linuxnamespaces/__init__.py b/linuxnamespaces/__init__.py
index b50f113..a39ab1d 100644
--- a/linuxnamespaces/__init__.py
+++ b/linuxnamespaces/__init__.py
@@ -136,6 +136,41 @@ class IDAllocation:
         """
         return IDMapping(target, self.allocate(count), count)
 
+    def reserve(self, start: int, count: int) -> None:
+        """Reserve (and remove) the given range from this allocation. If the
+        range is not fully contained in this allocation, a ValueError is
+        raised.
+        """
+        if count < 0:
+            raise ValueError("nagative count")
+        index = bisect.bisect_right(self.ranges, (start, float("inf"))) - 1
+        if index < 0:
+            raise ValueError("range to reserve not found")
+        cur_start, cur_count = self.ranges[index]
+        assert cur_start <= start
+        if cur_start == start:
+            # Requested range starts at range boundary
+            if cur_count < count:
+                raise ValueError("range to reserve not found")
+            if cur_count == count:
+                # Requested range matches a range exactly
+                del self.ranges[index]
+            else:
+                # Requested range is a head of the matched range
+                self.ranges[index] = (start + count, cur_count - count)
+        elif cur_start + cur_count >= start + count:
+            # Requested range fits into a matched range
+            self.ranges[index] = (cur_start, start - cur_start)
+            if cur_start + cur_count > start + count:
+                # Requested range punches a hole into a matched range
+                self.ranges.insert(
+                    index + 1,
+                    (start + count, cur_start + cur_count - (start + count))
+                )
+            # else: Requested range is a tail of a matched range
+        else:
+            raise ValueError("range to reserve not found")
+
 
 def newidmap(
     kind: typing.Literal["uid", "gid"],
diff --git a/tests/test_simple.py b/tests/test_simple.py
index 212f414..456e088 100644
--- a/tests/test_simple.py
+++ b/tests/test_simple.py
@@ -48,6 +48,24 @@ class IDAllocationTest(unittest.TestCase):
         alloc.add_range(3, 2)
         self.assertIn(alloc.allocate(3), (1, 2))
 
+    def test_reserve(self) -> None:
+        alloc = linuxnamespaces.IDAllocation()
+        alloc.add_range(0, 10)
+        # Split a range
+        alloc.reserve(3, 3)
+        self.assertEqual(alloc.ranges, [(0, 3), (6, 4)])
+        self.assertRaises(ValueError, alloc.reserve, 0, 4)
+        self.assertRaises(ValueError, alloc.reserve, 5, 4)
+        # Head of range
+        alloc.reserve(0, 2)
+        self.assertEqual(alloc.ranges, [(2, 1), (6, 4)])
+        # Tail of range
+        alloc.reserve(7, 3)
+        self.assertEqual(alloc.ranges, [(2, 1), (6, 1)])
+        # Exact range
+        alloc.reserve(6, 1)
+        self.assertEqual(alloc.ranges, [(2, 1)])
+
 
 class AsnycioTest(unittest.IsolatedAsyncioTestCase):
     async def test_eventfd(self) -> None:
-- 
cgit v1.2.3