src/allmydata/test/test_helper.py

   1
   2 import os
   3 from twisted.trial import unittest
   4 from twisted.application import service
   5
   6 from foolscap import Tub, eventual
   7 from foolscap.logging import log
   8
   9 from allmydata import offloaded, storage
  10 from allmydata.immutable import upload
  11 from allmydata.util import hashutil, fileutil, mathutil
  12 from pycryptopp.cipher.aes import AES
  13
  14 MiB = 1024*1024
  15
  16 DATA = "I need help\n" * 1000
  17
  18 class CHKUploadHelper_fake(offloaded.CHKUploadHelper):
  19     def start_encrypted(self, eu):
  20         d = eu.get_size()
  21         def _got_size(size):
  22             d2 = eu.get_all_encoding_parameters()
  23             def _got_parms(parms):
  24                 needed_shares, happy, total_shares, segsize = parms
  25                 ueb_data = {"needed_shares": needed_shares,
  26                             "total_shares": total_shares,
  27                             "segment_size": segsize,
  28                             "size": size,
  29                             }
  30                 self._results.uri_extension_data = ueb_data
  31                 return (hashutil.uri_extension_hash(""),
  32                         needed_shares, total_shares, size)
  33             d2.addCallback(_got_parms)
  34             return d2
  35         d.addCallback(_got_size)
  36         return d
  37
  38 class CHKUploadHelper_already_uploaded(offloaded.CHKUploadHelper):
  39     def start(self):
  40         res = upload.UploadResults()
  41         res.uri_extension_hash = hashutil.uri_extension_hash("")
  42
  43         # we're pretending that the file they're trying to upload was already
  44         # present in the grid. We return some information about the file, so
  45         # the client can decide if they like the way it looks. The parameters
  46         # used here are chosen to match the defaults.
  47         PARAMS = FakeClient.DEFAULT_ENCODING_PARAMETERS
  48         ueb_data = {"needed_shares": PARAMS["k"],
  49                     "total_shares": PARAMS["n"],
  50                     "segment_size": min(PARAMS["max_segment_size"], len(DATA)),
  51                     "size": len(DATA),
  52                     }
  53         res.uri_extension_data = ueb_data
  54         return (res, None)
  55
  56 class FakeClient(service.MultiService):
  57     DEFAULT_ENCODING_PARAMETERS = {"k":25,
  58                                    "happy": 75,
  59                                    "n": 100,
  60                                    "max_segment_size": 1*MiB,
  61                                    }
  62     stats_provider = None
  63     def log(self, *args, **kwargs):
  64         return log.msg(*args, **kwargs)
  65     def get_encoding_parameters(self):
  66         return self.DEFAULT_ENCODING_PARAMETERS
  67     def get_permuted_peers(self, service_name, storage_index):
  68         return []
  69
  70 def flush_but_dont_ignore(res):
  71     d = eventual.flushEventualQueue()
  72     def _done(ignored):
  73         return res
  74     d.addCallback(_done)
  75     return d
  76
  77 def wait_a_few_turns(ignored=None):
  78     d = eventual.fireEventually()
  79     d.addCallback(eventual.fireEventually)
  80     d.addCallback(eventual.fireEventually)
  81     d.addCallback(eventual.fireEventually)
  82     d.addCallback(eventual.fireEventually)
  83     d.addCallback(eventual.fireEventually)
  84     return d
  85
  86 def upload_data(uploader, data, convergence):
  87     u = upload.Data(data, convergence=convergence)
  88     return uploader.upload(u)
  89
  90 class AssistedUpload(unittest.TestCase):
  91     def setUp(self):
  92         self.s = FakeClient()
  93         self.s.startService()
  94
  95         self.tub = t = Tub()
  96         t.setServiceParent(self.s)
  97         self.s.tub = t
  98         # we never actually use this for network traffic, so it can use a
  99         # bogus host/port
 100         t.setLocation("bogus:1234")
 101
 102     def setUpHelper(self, basedir):
 103         fileutil.make_dirs(basedir)
 104         self.helper = h = offloaded.Helper(basedir)
 105         h.chk_upload_helper_class = CHKUploadHelper_fake
 106         h.setServiceParent(self.s)
 107         self.helper_furl = self.tub.registerReference(h)
 108
 109     def tearDown(self):
 110         d = self.s.stopService()
 111         d.addCallback(eventual.fireEventually)
 112         d.addBoth(flush_but_dont_ignore)
 113         return d
 114
 115
 116     def test_one(self):
 117         self.basedir = "helper/AssistedUpload/test_one"
 118         self.setUpHelper(self.basedir)
 119         u = upload.Uploader(self.helper_furl)
 120         u.setServiceParent(self.s)
 121
 122         d = wait_a_few_turns()
 123
 124         def _ready(res):
 125             assert u._helper
 126
 127             return upload_data(u, DATA, convergence="some convergence string")
 128         d.addCallback(_ready)
 129         def _uploaded(results):
 130             uri = results.uri
 131             assert "CHK" in uri
 132         d.addCallback(_uploaded)
 133
 134         def _check_empty(res):
 135             files = os.listdir(os.path.join(self.basedir, "CHK_encoding"))
 136             self.failUnlessEqual(files, [])
 137             files = os.listdir(os.path.join(self.basedir, "CHK_incoming"))
 138             self.failUnlessEqual(files, [])
 139         d.addCallback(_check_empty)
 140
 141         return d
 142
 143     def test_previous_upload_failed(self):
 144         self.basedir = "helper/AssistedUpload/test_previous_upload_failed"
 145         self.setUpHelper(self.basedir)
 146
 147         # we want to make sure that an upload which fails (leaving the
 148         # ciphertext in the CHK_encoding/ directory) does not prevent a later
 149         # attempt to upload that file from working. We simulate this by
 150         # populating the directory manually. The hardest part is guessing the
 151         # storage index.
 152
 153         k = FakeClient.DEFAULT_ENCODING_PARAMETERS["k"]
 154         n = FakeClient.DEFAULT_ENCODING_PARAMETERS["n"]
 155         max_segsize = FakeClient.DEFAULT_ENCODING_PARAMETERS["max_segment_size"]
 156         segsize = min(max_segsize, len(DATA))
 157         # this must be a multiple of 'required_shares'==k
 158         segsize = mathutil.next_multiple(segsize, k)
 159
 160         key = hashutil.convergence_hash(k, n, segsize, DATA, "test convergence string")
 161         assert len(key) == 16
 162         encryptor = AES(key)
 163         SI = hashutil.storage_index_hash(key)
 164         SI_s = storage.si_b2a(SI)
 165         encfile = os.path.join(self.basedir, "CHK_encoding", SI_s)
 166         f = open(encfile, "wb")
 167         f.write(encryptor.process(DATA))
 168         f.close()
 169
 170         u = upload.Uploader(self.helper_furl)
 171         u.setServiceParent(self.s)
 172
 173         d = wait_a_few_turns()
 174
 175         def _ready(res):
 176             assert u._helper
 177             return upload_data(u, DATA, convergence="test convergence string")
 178         d.addCallback(_ready)
 179         def _uploaded(results):
 180             uri = results.uri
 181             assert "CHK" in uri
 182         d.addCallback(_uploaded)
 183
 184         def _check_empty(res):
 185             files = os.listdir(os.path.join(self.basedir, "CHK_encoding"))
 186             self.failUnlessEqual(files, [])
 187             files = os.listdir(os.path.join(self.basedir, "CHK_incoming"))
 188             self.failUnlessEqual(files, [])
 189         d.addCallback(_check_empty)
 190
 191         return d
 192
 193     def test_already_uploaded(self):
 194         self.basedir = "helper/AssistedUpload/test_already_uploaded"
 195         self.setUpHelper(self.basedir)
 196         self.helper.chk_upload_helper_class = CHKUploadHelper_already_uploaded
 197         u = upload.Uploader(self.helper_furl)
 198         u.setServiceParent(self.s)
 199
 200         d = wait_a_few_turns()
 201
 202         def _ready(res):
 203             assert u._helper
 204
 205             return upload_data(u, DATA, convergence="some convergence string")
 206         d.addCallback(_ready)
 207         def _uploaded(results):
 208             uri = results.uri
 209             assert "CHK" in uri
 210         d.addCallback(_uploaded)
 211
 212         def _check_empty(res):
 213             files = os.listdir(os.path.join(self.basedir, "CHK_encoding"))
 214             self.failUnlessEqual(files, [])
 215             files = os.listdir(os.path.join(self.basedir, "CHK_incoming"))
 216             self.failUnlessEqual(files, [])
 217         d.addCallback(_check_empty)
 218
 219         return d