From 49dfc8445cec28d6d903d0a15ee69c411b1b70a1 Mon Sep 17 00:00:00 2001
From: Itamar Turner-Trauring <itamar@pythonspeed.com>
Date: Mon, 18 Jul 2022 14:12:12 -0400
Subject: [PATCH] Implementation of getting length of shares (albeit
 inefficiently for now).

---
 src/allmydata/storage/immutable.py |  8 ++++++++
 src/allmydata/storage/mutable.py   | 10 +++++-----
 src/allmydata/storage/server.py    | 10 ++++++++++
 src/allmydata/test/test_storage.py | 22 ++++++++++++++++++++++
 4 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/src/allmydata/storage/immutable.py b/src/allmydata/storage/immutable.py
index 920bd3c5e..2c65304b8 100644
--- a/src/allmydata/storage/immutable.py
+++ b/src/allmydata/storage/immutable.py
@@ -199,8 +199,13 @@ class ShareFile(object):
                 raise UnknownImmutableContainerVersionError(filename, version)
             self._num_leases = num_leases
             self._lease_offset = filesize - (num_leases * self.LEASE_SIZE)
+            self._length = filesize - 0xc - (num_leases * self.LEASE_SIZE)
+
         self._data_offset = 0xc
 
+    def get_length(self):
+        return self._length
+
     def unlink(self):
         os.unlink(self.home)
 
@@ -544,6 +549,9 @@ class BucketReader(object):
                                             self.shnum,
                                             reason)
 
+    def get_length(self):
+        return self._share_file.get_length()
+
 
 @implementer(RIBucketReader)
 class FoolscapBucketReader(Referenceable):  # type: ignore # warner/foolscap#78
diff --git a/src/allmydata/storage/mutable.py b/src/allmydata/storage/mutable.py
index bd59d96b8..9a99979e9 100644
--- a/src/allmydata/storage/mutable.py
+++ b/src/allmydata/storage/mutable.py
@@ -412,11 +412,11 @@ class MutableShareFile(object):
                 datav.append(self._read_share_data(f, offset, length))
         return datav
 
-#    def remote_get_length(self):
-#        f = open(self.home, 'rb')
-#        data_length = self._read_data_length(f)
-#        f.close()
-#        return data_length
+    def get_length(self):
+        f = open(self.home, 'rb')
+        data_length = self._read_data_length(f)
+        f.close()
+        return data_length
 
     def check_write_enabler(self, write_enabler, si_s):
         with open(self.home, 'rb+') as f:
diff --git a/src/allmydata/storage/server.py b/src/allmydata/storage/server.py
index 0a1999dfb..f452885d0 100644
--- a/src/allmydata/storage/server.py
+++ b/src/allmydata/storage/server.py
@@ -794,6 +794,16 @@ class StorageServer(service.MultiService):
 
         return None
 
+    def get_immutable_share_length(self, storage_index: bytes, share_number: int) -> int:
+        """Returns the length (in bytes) of an immutable."""
+        return self.get_buckets(storage_index)[share_number].get_length()
+
+    def get_mutable_share_length(self, storage_index: bytes, share_number: int) -> int:
+        """Returns the length (in bytes) of a mutable."""
+        return MutableShareFile(
+            dict(self.get_shares(storage_index))[share_number]
+        ).get_length()
+
 
 @implementer(RIStorageServer)
 class FoolscapStorageServer(Referenceable):  # type: ignore # warner/foolscap#78
diff --git a/src/allmydata/test/test_storage.py b/src/allmydata/test/test_storage.py
index 91d55790e..bb8d48d2f 100644
--- a/src/allmydata/test/test_storage.py
+++ b/src/allmydata/test/test_storage.py
@@ -688,6 +688,15 @@ class Server(unittest.TestCase):
             writer.abort()
         self.failUnlessEqual(ss.allocated_size(), 0)
 
+    def test_immutable_length(self):
+        """``get_immutable_share_length()`` returns the length of an immutable share."""
+        ss = self.create("test_immutable_length")
+        _, writers = self.allocate(ss, b"allocate", [22], 75)
+        bucket = writers[22]
+        bucket.write(0, b"X" * 75)
+        bucket.close()
+        self.assertEqual(ss.get_immutable_share_length(b"allocate", 22), 75)
+
     def test_allocate(self):
         ss = self.create("test_allocate")
 
@@ -1340,6 +1349,19 @@ class MutableServer(unittest.TestCase):
             (set(), {0, 1, 2, 4}, {0, 1, 4})
         )
 
+    def test_mutable_share_length(self):
+        """``get_mutable_share_length()`` returns the length of the share."""
+        ss = self.create("test_mutable_share_length")
+        self.allocate(ss, b"si1", b"we1", b"le1", [16], 23)
+        ss.slot_testv_and_readv_and_writev(
+            b"si1", (self.write_enabler(b"we1"),
+                     self.renew_secret(b"le1"),
+                     self.cancel_secret(b"le1")),
+            {16: ([], [(0, b"x" * 23)], None)},
+            []
+        )
+        self.assertEqual(ss.get_mutable_share_length(b"si1", 16), 23)
+
     def test_bad_magic(self):
         ss = self.create("test_bad_magic")
         self.allocate(ss, b"si1", b"we1", next(self._lease_secret), set([0]), 10)