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 upload, offloaded
  10 from allmydata.util import hashutil, fileutil, idlib
  11 from pycryptopp.cipher.aes import AES
  12
  13 MiB = 1024*1024
  14
  15 class CHKUploadHelper_fake(offloaded.CHKUploadHelper):
  16     def start_encrypted(self, eu):
  17         d = eu.get_size()
  18         def _got_size(size):
  19             d2 = eu.get_all_encoding_parameters()
  20             def _got_parms(parms):
  21                 needed_shares, happy, total_shares, segsize = parms
  22                 return (hashutil.uri_extension_hash(""),
  23                         needed_shares, total_shares, size)
  24             d2.addCallback(_got_parms)
  25             return d2
  26         d.addCallback(_got_size)
  27         return d
  28
  29 class CHKUploadHelper_already_uploaded(offloaded.CHKUploadHelper):
  30     def start(self):
  31         res = {'uri_extension_hash': hashutil.uri_extension_hash("")}
  32         return (res, None)
  33
  34 class FakeClient(service.MultiService):
  35     DEFAULT_ENCODING_PARAMETERS = {"k":25,
  36                                    "happy": 75,
  37                                    "n": 100,
  38                                    "max_segment_size": 1*MiB,
  39                                    }
  40     def log(self, *args, **kwargs):
  41         return log.msg(*args, **kwargs)
  42     def get_push_to_ourselves(self):
  43         return True
  44     def get_encoding_parameters(self):
  45         return self.DEFAULT_ENCODING_PARAMETERS
  46     def get_permuted_peers(self, storage_index):
  47         return []
  48
  49 def flush_but_dont_ignore(res):
  50     d = eventual.flushEventualQueue()
  51     def _done(ignored):
  52         return res
  53     d.addCallback(_done)
  54     return d
  55
  56 def upload_data(uploader, data):
  57     u = upload.Data(data)
  58     return uploader.upload(u)
  59
  60 class AssistedUpload(unittest.TestCase):
  61     def setUp(self):
  62         self.s = FakeClient()
  63         self.s.startService()
  64
  65         self.tub = t = Tub()
  66         t.setServiceParent(self.s)
  67         self.s.tub = t
  68         # we never actually use this for network traffic, so it can use a
  69         # bogus host/port
  70         t.setLocation("bogus:1234")
  71
  72     def setUpHelper(self, basedir):
  73         fileutil.make_dirs(basedir)
  74         self.helper = h = offloaded.Helper(basedir)
  75         h.chk_upload_helper_class = CHKUploadHelper_fake
  76         h.setServiceParent(self.s)
  77         self.helper_furl = self.tub.registerReference(h)
  78
  79     def tearDown(self):
  80         d = self.s.stopService()
  81         d.addCallback(eventual.fireEventually)
  82         d.addBoth(flush_but_dont_ignore)
  83         return d
  84
  85
  86     def test_one(self):
  87         self.basedir = "helper/AssistedUpload/test_one"
  88         self.setUpHelper(self.basedir)
  89         u = upload.Uploader(self.helper_furl)
  90         u.setServiceParent(self.s)
  91
  92         # wait a few turns
  93         d = eventual.fireEventually()
  94         d.addCallback(eventual.fireEventually)
  95         d.addCallback(eventual.fireEventually)
  96
  97         def _ready(res):
  98             assert u._helper
  99
 100             DATA = "I need help\n" * 1000
 101             return upload_data(u, DATA)
 102         d.addCallback(_ready)
 103         def _uploaded(uri):
 104             assert "CHK" in uri
 105         d.addCallback(_uploaded)
 106
 107         def _check_empty(res):
 108             files = os.listdir(os.path.join(self.basedir, "CHK_encoding"))
 109             self.failUnlessEqual(files, [])
 110             files = os.listdir(os.path.join(self.basedir, "CHK_incoming"))
 111             self.failUnlessEqual(files, [])
 112         d.addCallback(_check_empty)
 113
 114         return d
 115
 116     def test_previous_upload_failed(self):
 117         self.basedir = "helper/AssistedUpload/test_previous_upload_failed"
 118         self.setUpHelper(self.basedir)
 119         DATA = "I need help\n" * 1000
 120
 121         # we want to make sure that an upload which fails (leaving the
 122         # ciphertext in the CHK_encoding/ directory) does not prevent a later
 123         # attempt to upload that file from working. We simulate this by
 124         # populating the directory manually.
 125         key = hashutil.key_hash(DATA)[:16]
 126         encryptor = AES(key)
 127         SI = hashutil.storage_index_hash(key)
 128         SI_s = idlib.b2a(SI)
 129         encfile = os.path.join(self.basedir, "CHK_encoding", SI_s)
 130         f = open(encfile, "wb")
 131         f.write(encryptor.process(DATA))
 132         f.close()
 133
 134         u = upload.Uploader(self.helper_furl)
 135         u.setServiceParent(self.s)
 136
 137         # wait a few turns
 138         d = eventual.fireEventually()
 139         d.addCallback(eventual.fireEventually)
 140         d.addCallback(eventual.fireEventually)
 141
 142         def _ready(res):
 143             assert u._helper
 144             return upload_data(u, DATA)
 145         d.addCallback(_ready)
 146         def _uploaded(uri):
 147             assert "CHK" in uri
 148         d.addCallback(_uploaded)
 149
 150         def _check_empty(res):
 151             files = os.listdir(os.path.join(self.basedir, "CHK_encoding"))
 152             self.failUnlessEqual(files, [])
 153             files = os.listdir(os.path.join(self.basedir, "CHK_incoming"))
 154             self.failUnlessEqual(files, [])
 155         d.addCallback(_check_empty)
 156
 157         return d
 158
 159     def test_already_uploaded(self):
 160         self.basedir = "helper/AssistedUpload/test_already_uploaded"
 161         self.setUpHelper(self.basedir)
 162         self.helper.chk_upload_helper_class = CHKUploadHelper_already_uploaded
 163         u = upload.Uploader(self.helper_furl)
 164         u.setServiceParent(self.s)
 165
 166         # wait a few turns
 167         d = eventual.fireEventually()
 168         d.addCallback(eventual.fireEventually)
 169         d.addCallback(eventual.fireEventually)
 170
 171         def _ready(res):
 172             assert u._helper
 173
 174             DATA = "I need help\n" * 1000
 175             return upload_data(u, DATA)
 176         d.addCallback(_ready)
 177         def _uploaded(uri):
 178             assert "CHK" in uri
 179         d.addCallback(_uploaded)
 180
 181         def _check_empty(res):
 182             files = os.listdir(os.path.join(self.basedir, "CHK_encoding"))
 183             self.failUnlessEqual(files, [])
 184             files = os.listdir(os.path.join(self.basedir, "CHK_incoming"))
 185             self.failUnlessEqual(files, [])
 186         d.addCallback(_check_empty)
 187
 188         return d