diff options
| -rw-r--r-- | linuxnamespaces/__init__.py | 35 | ||||
| -rw-r--r-- | tests/test_simple.py | 18 |
2 files changed, 53 insertions, 0 deletions
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: |