From 94e6e6160b24959c28322f659d984e6d692e418b Mon Sep 17 00:00:00 2001
From: Brian Warner <warner@allmydata.com>
Date: Tue, 3 Jul 2007 17:27:07 -0700
Subject: [PATCH] activate storage size limits in the client. Closes #34.

To use this, write a number like 10MB or 5Gb or 5000000000 to a file
named 'sizelimit' in the client's base directory. The node will not grant
leases for shares that would take it much beyond this many bytes of
storage. Note that metadata is not included in the allocation count until
a restart, so the actual space consumed may grow beyond the limit if
the node is not restarted very frequently and the amount of metadata is
significant.
---
 src/allmydata/client.py           | 31 +++++++++++++++++--
 src/allmydata/test/test_client.py | 50 +++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 2 deletions(-)

diff --git a/src/allmydata/client.py b/src/allmydata/client.py
index 80a0bbaa..fc7b092b 100644
--- a/src/allmydata/client.py
+++ b/src/allmydata/client.py
@@ -1,5 +1,5 @@
 
-import os, sha, stat, time
+import os, sha, stat, time, re
 from foolscap import Referenceable, SturdyRef
 from zope.interface import implements
 from allmydata.interfaces import RIClient
@@ -7,6 +7,7 @@ from allmydata import node
 
 from twisted.internet import reactor
 from twisted.application.internet import TimerService
+from twisted.python import log
 
 import allmydata
 from allmydata.Crypto.Util.number import bytes_to_long
@@ -27,6 +28,7 @@ class Client(node.Node, Referenceable):
     INTRODUCER_FURL_FILE = "introducer.furl"
     MY_FURL_FILE = "myself.furl"
     SUICIDE_PREVENTION_HOTLINE_FILE = "suicide_prevention_hotline"
+    SIZELIMIT_FILE = "sizelimit"
 
     # we're pretty narrow-minded right now
     OLDEST_SUPPORTED_VERSION = allmydata.__version__
@@ -35,7 +37,7 @@ class Client(node.Node, Referenceable):
         node.Node.__init__(self, basedir)
         self.my_furl = None
         self.introducer_client = None
-        self.add_service(StorageServer(os.path.join(basedir, self.STOREDIR)))
+        self.init_storage()
         self.add_service(Uploader())
         self.add_service(Downloader())
         self.add_service(VirtualDrive())
@@ -58,6 +60,31 @@ class Client(node.Node, Referenceable):
             hotline = TimerService(1.0, self._check_hotline, hotline_file)
             hotline.setServiceParent(self)
 
+    def init_storage(self):
+        storedir = os.path.join(self.basedir, self.STOREDIR)
+        sizelimit = None
+        SIZELIMIT_FILE = os.path.join(self.basedir,
+                                      self.SIZELIMIT_FILE)
+        if os.path.exists(SIZELIMIT_FILE):
+            f = open(SIZELIMIT_FILE, "r")
+            data = f.read().strip()
+            f.close()
+            m = re.match(r"^(\d+)([kKmMgG]?[bB]?)$", data)
+            if not m:
+                log.msg("SIZELIMIT_FILE contains unparseable value %s" % data)
+            else:
+                number, suffix = m.groups()
+                suffix = suffix.upper()
+                if suffix.endswith("B"):
+                    suffix = suffix[:-1]
+                multiplier = {"": 1,
+                              "K": 1000,
+                              "M": 1000 * 1000,
+                              "G": 1000 * 1000 * 1000,
+                              }[suffix]
+                sizelimit = int(number) * multiplier
+        self.add_service(StorageServer(storedir, sizelimit))
+
     def _check_hotline(self, hotline_file):
         if os.path.exists(hotline_file):
             mtime = os.stat(hotline_file)[stat.ST_MTIME]
diff --git a/src/allmydata/test/test_client.py b/src/allmydata/test/test_client.py
index 167454cc..f55980ec 100644
--- a/src/allmydata/test/test_client.py
+++ b/src/allmydata/test/test_client.py
@@ -29,6 +29,56 @@ class Basic(unittest.TestCase):
         open(os.path.join(basedir, "introducer.furl"), "w").write("")
         c = client.Client(basedir)
 
+    def test_sizelimit_1(self):
+        basedir = "client.Basic.test_sizelimit_1"
+        os.mkdir(basedir)
+        open(os.path.join(basedir, "introducer.furl"), "w").write("")
+        open(os.path.join(basedir, "vdrive.furl"), "w").write("")
+        open(os.path.join(basedir, "sizelimit"), "w").write("1000")
+        c = client.Client(basedir)
+        self.failUnlessEqual(c.getServiceNamed("storageserver").sizelimit,
+                             1000)
+
+    def test_sizelimit_2(self):
+        basedir = "client.Basic.test_sizelimit_2"
+        os.mkdir(basedir)
+        open(os.path.join(basedir, "introducer.furl"), "w").write("")
+        open(os.path.join(basedir, "vdrive.furl"), "w").write("")
+        open(os.path.join(basedir, "sizelimit"), "w").write("10K")
+        c = client.Client(basedir)
+        self.failUnlessEqual(c.getServiceNamed("storageserver").sizelimit,
+                             10*1000)
+
+    def test_sizelimit_3(self):
+        basedir = "client.Basic.test_sizelimit_3"
+        os.mkdir(basedir)
+        open(os.path.join(basedir, "introducer.furl"), "w").write("")
+        open(os.path.join(basedir, "vdrive.furl"), "w").write("")
+        open(os.path.join(basedir, "sizelimit"), "w").write("5mB")
+        c = client.Client(basedir)
+        self.failUnlessEqual(c.getServiceNamed("storageserver").sizelimit,
+                             5*1000*1000)
+
+    def test_sizelimit_4(self):
+        basedir = "client.Basic.test_sizelimit_4"
+        os.mkdir(basedir)
+        open(os.path.join(basedir, "introducer.furl"), "w").write("")
+        open(os.path.join(basedir, "vdrive.furl"), "w").write("")
+        open(os.path.join(basedir, "sizelimit"), "w").write("78Gb")
+        c = client.Client(basedir)
+        self.failUnlessEqual(c.getServiceNamed("storageserver").sizelimit,
+                             78*1000*1000*1000)
+
+    def test_sizelimit_bad(self):
+        basedir = "client.Basic.test_sizelimit_bad"
+        os.mkdir(basedir)
+        open(os.path.join(basedir, "introducer.furl"), "w").write("")
+        open(os.path.join(basedir, "vdrive.furl"), "w").write("")
+        open(os.path.join(basedir, "sizelimit"), "w").write("bogus")
+        c = client.Client(basedir)
+        self.failUnlessEqual(c.getServiceNamed("storageserver").sizelimit,
+                             None)
+
     def test_permute(self):
         basedir = "test_client.Basic.test_permute"
         os.mkdir(basedir)
-- 
2.45.2