make it possible to test this behavior of add_lease

src/allmydata/storage/immutable.py

@@ -53,13 +53,64 @@ from allmydata.storage.common import UnknownImmutableContainerVersionError
 # then the value stored in this field will be the actual share data length
 # modulo 2**32.
 
+def _fix_lease_count_format(lease_count_format):
+    """
+    Turn a single character struct format string into a format string suitable
+    for use in encoding and decoding the lease count value inside a share
+    file, if possible.
+
+    :param str lease_count_format: A single character format string like
+        ``"B"`` or ``"L"``.
+
+    :raise ValueError: If the given format string is not suitable for use
+        encoding and decoding a lease count.
+
+    :return str: A complete format string which can safely be used to encode
+        and decode lease counts in a share file.
+    """
+    if len(lease_count_format) != 1:
+        raise ValueError(
+            "Cannot construct ShareFile with lease_count_format={!r}; "
+            "format must accept a single value".format(
+                lease_count_format,
+            ),
+        )
+    # Make it big-endian with standard size so all platforms agree on the
+    # result.
+    fixed = ">" + lease_count_format
+    if struct.calcsize(fixed) > 4:
+        # There is only room for at most 4 bytes in the share file format so
+        # we can't allow any larger formats.
+        raise ValueError(
+            "Cannot construct ShareFile with lease_count_format={!r}; "
+            "size must be smaller than size of '>L'".format(
+                lease_count_format,
+            ),
+        )
+    return fixed
+
+
 class ShareFile(object):
+    """
+    Support interaction with persistent storage of a share.
+
+    :ivar str _lease_count_format: The format string which is used to encode
+        and decode the lease count inside the share file.  As stated in the
+        comment in this module there is room for at most 4 bytes in this part
+        of the file.  A format string that works on fewer bytes is allowed to
+        restrict the number of leases allowed in the share file to a smaller
+        number than could be supported by using the full 4 bytes.  This is
+        mostly of interest for testing.
+    """
     LEASE_SIZE = struct.calcsize(">L32s32sL")
     sharetype = "immutable"
 
-    def __init__(self, filename, max_size=None, create=False):
+    def __init__(self, filename, max_size=None, create=False, lease_count_format="L"):
         """ If max_size is not None then I won't allow more than max_size to be written to me. If create=True and max_size must not be None. """
         precondition((max_size is not None) or (not create), max_size, create)
+
+        self._lease_count_format = _fix_lease_count_format(lease_count_format)
+        self._lease_count_size = struct.calcsize(self._lease_count_format)
         self.home = filename
         self._max_size = max_size
         if create:
@@ -126,12 +177,21 @@ class ShareFile(object):
 
     def _read_num_leases(self, f):
         f.seek(0x08)
-        (num_leases,) = struct.unpack(">L", f.read(4))
+        (num_leases,) = struct.unpack(
+            self._lease_count_format,
+            f.read(self._lease_count_size),
+        )
         return num_leases
 
     def _write_num_leases(self, f, num_leases):
+        self._write_encoded_num_leases(
+            f,
+            struct.pack(self._lease_count_format, num_leases),
+        )
+
+    def _write_encoded_num_leases(self, f, encoded_num_leases):
         f.seek(0x08)
-        f.write(struct.pack(">L", num_leases))
+        f.write(encoded_num_leases)
 
     def _truncate_leases(self, f, num_leases):
         f.truncate(self._lease_offset + num_leases * self.LEASE_SIZE)

src/allmydata/test/test_storage.py

@@ -19,6 +19,7 @@ import platform
 import stat
 import struct
 import shutil
+from functools import partial
 from uuid import uuid4
 
 from twisted.trial import unittest
@@ -3009,8 +3010,8 @@ class Stats(unittest.TestCase):
 class ShareFileTests(unittest.TestCase):
     """Tests for allmydata.storage.immutable.ShareFile."""
 
-    def get_sharefile(self):
-        sf = ShareFile(self.mktemp(), max_size=1000, create=True)
+    def get_sharefile(self, **kwargs):
+        sf = ShareFile(self.mktemp(), max_size=1000, create=True, **kwargs)
         sf.write_share_data(0, b"abc")
         sf.write_share_data(2, b"DEF")
         # Should be b'abDEF' now.
@@ -3039,3 +3040,57 @@ class ShareFileTests(unittest.TestCase):
         sf = self.get_sharefile()
         with self.assertRaises(IndexError):
             sf.cancel_lease(b"garbage")
+
+    def test_long_lease_count_format(self):
+        """
+        ``ShareFile.__init__`` raises ``ValueError`` if the lease count format
+        given is longer than one character.
+        """
+        with self.assertRaises(ValueError):
+            self.get_sharefile(lease_count_format="BB")
+
+    def test_large_lease_count_format(self):
+        """
+        ``ShareFile.__init__`` raises ``ValueError`` if the lease count format
+        encodes to a size larger than 8 bytes.
+        """
+        with self.assertRaises(ValueError):
+            self.get_sharefile(lease_count_format="Q")
+
+    def test_avoid_lease_overflow(self):
+        """
+        If the share file already has the maximum number of leases supported then
+        ``ShareFile.add_lease`` raises ``struct.error`` and makes no changes
+        to the share file contents.
+        """
+        make_lease = partial(
+            LeaseInfo,
+            renew_secret=b"r" * 32,
+            cancel_secret=b"c" * 32,
+            expiration_time=2 ** 31,
+        )
+        # Make it a little easier to reach the condition by limiting the
+        # number of leases to only 255.
+        sf = self.get_sharefile(lease_count_format="B")
+
+        # Add the leases.
+        for i in range(2 ** 8 - 1):
+            lease = make_lease(owner_num=i)
+            sf.add_lease(lease)
+
+        # Capture the state of the share file at this point so we can
+        # determine whether the next operation modifies it or not.
+        with open(sf.home, "rb") as f:
+            before_data = f.read()
+
+        # It is not possible to add a 256th lease.
+        lease = make_lease(owner_num=256)
+        with self.assertRaises(struct.error):
+            sf.add_lease(lease)
+
+        # Compare the share file state to what we captured earlier.  Any
+        # change is a bug.
+        with open(sf.home, "rb") as f:
+            after_data = f.read()
+
+        self.assertEqual(before_data, after_data)