PyORAM/src/pyoram/tests/test_heap_storage.py

   1 import os
   2 import unittest2
   3 import tempfile
   4
   5 from pyoram.util.virtual_heap import \
   6     SizedVirtualHeap
   7 from pyoram.storage.block_storage import \
   8     BlockStorageTypeFactory
   9 from pyoram.storage.block_storage_file import \
  10     BlockStorageFile
  11 from pyoram.storage.heap_storage import \
  12     HeapStorage
  13
  14 from six.moves import xrange
  15
  16 thisdir = os.path.dirname(os.path.abspath(__file__))
  17
  18 class TestHeapStorage(unittest2.TestCase):
  19
  20     @classmethod
  21     def setUpClass(cls):
  22         fd, cls._dummy_name = tempfile.mkstemp()
  23         os.close(fd)
  24         try:
  25             os.remove(cls._dummy_name)
  26         except OSError:                                # pragma: no cover
  27             pass                                       # pragma: no cover
  28         cls._block_size = 25
  29         cls._blocks_per_bucket = 3
  30         cls._heap_base = 4
  31         cls._heap_height = 2
  32         cls._bucket_count = \
  33             ((cls._heap_base**(cls._heap_height+1)) - 1)//(cls._heap_base-1)
  34         cls._block_count = cls._bucket_count * \
  35                            cls._blocks_per_bucket
  36         cls._testfname = cls.__name__ + "_testfile.bin"
  37         cls._buckets = []
  38         cls._type_name = "file"
  39         f = HeapStorage.setup(
  40             cls._testfname,
  41             block_size=cls._block_size,
  42             heap_height=cls._heap_height,
  43             heap_base=cls._heap_base,
  44             blocks_per_bucket=cls._blocks_per_bucket,
  45             storage_type=cls._type_name,
  46             initialize=lambda i: bytes(bytearray([i]) * \
  47                                        cls._block_size * \
  48                                        cls._blocks_per_bucket),
  49             ignore_existing=True)
  50         f.close()
  51         for i in range(cls._bucket_count):
  52             data = bytearray([i]) * \
  53                    cls._block_size * \
  54                    cls._blocks_per_bucket
  55             cls._buckets.append(data)
  56
  57     @classmethod
  58     def tearDownClass(cls):
  59         try:
  60             os.remove(cls._testfname)
  61         except OSError:                                # pragma: no cover
  62             pass                                       # pragma: no cover
  63         try:
  64             os.remove(cls._dummy_name)
  65         except OSError:                                # pragma: no cover
  66             pass                                       # pragma: no cover
  67
  68     def test_setup_fails(self):
  69         self.assertEqual(os.path.exists(self._dummy_name), False)
  70         with self.assertRaises(IOError):
  71             HeapStorage.setup(
  72                 os.path.join(thisdir,
  73                              "baselines",
  74                              "exists.empty"),
  75                 block_size=10,
  76                 heap_height=1,
  77                 blocks_per_bucket=1,
  78                 storage_type=self._type_name)
  79         self.assertEqual(os.path.exists(self._dummy_name), False)
  80         with self.assertRaises(IOError):
  81             HeapStorage.setup(
  82                 os.path.join(thisdir,
  83                              "baselines",
  84                              "exists.empty"),
  85                 block_size=10,
  86                 heap_height=1,
  87                 blocks_per_bucket=1,
  88                 storage_type=self._type_name,
  89                 ignore_existing=False)
  90         self.assertEqual(os.path.exists(self._dummy_name), False)
  91         # bad block_size
  92         with self.assertRaises(ValueError):
  93             HeapStorage.setup(
  94                 self._dummy_name,
  95                 block_size=0,
  96                 heap_height=1,
  97                 blocks_per_bucket=1,
  98                 storage_type=self._type_name)
  99         self.assertEqual(os.path.exists(self._dummy_name), False)
 100         # bad heap_height
 101         with self.assertRaises(ValueError):
 102             HeapStorage.setup(
 103                 self._dummy_name,
 104                 block_size=1,
 105                 heap_height=-1,
 106                 blocks_per_bucket=1,
 107                 storage_type=self._type_name)
 108         self.assertEqual(os.path.exists(self._dummy_name), False)
 109         # bad blocks_per_bucket
 110         with self.assertRaises(ValueError):
 111             HeapStorage.setup(
 112                 self._dummy_name,
 113                 block_size=1,
 114                 heap_height=1,
 115                 blocks_per_bucket=0,
 116                 storage_type=self._type_name)
 117         self.assertEqual(os.path.exists(self._dummy_name), False)
 118         # bad heap_base
 119         with self.assertRaises(ValueError):
 120             HeapStorage.setup(
 121                 self._dummy_name,
 122                 block_size=1,
 123                 heap_height=1,
 124                 blocks_per_bucket=1,
 125                 heap_base=1,
 126                 storage_type=self._type_name)
 127         self.assertEqual(os.path.exists(self._dummy_name), False)
 128         # bad header_data
 129         with self.assertRaises(TypeError):
 130             HeapStorage.setup(
 131                 self._dummy_name,
 132                 block_size=1,
 133                 heap_height=1,
 134                 blocks_per_bucket=1,
 135                 storage_type=self._type_name,
 136                 header_data=2)
 137         self.assertEqual(os.path.exists(self._dummy_name), False)
 138         # uses block_count
 139         with self.assertRaises(ValueError):
 140             HeapStorage.setup(
 141                 self._dummy_name,
 142                 block_size=1,
 143                 heap_height=1,
 144                 blocks_per_bucket=1,
 145                 block_count=1,
 146                 storage_type=self._type_name)
 147         self.assertEqual(os.path.exists(self._dummy_name), False)
 148
 149     def test_setup(self):
 150         fname = ".".join(self.id().split(".")[1:])
 151         fname += ".bin"
 152         fname = os.path.join(thisdir, fname)
 153         if os.path.exists(fname):
 154             os.remove(fname)                           # pragma: no cover
 155         bsize = 10
 156         heap_height = 2
 157         blocks_per_bucket = 3
 158         fsetup = HeapStorage.setup(
 159             fname,
 160             bsize,
 161             heap_height,
 162             blocks_per_bucket=blocks_per_bucket)
 163         fsetup.close()
 164         self.assertEqual(type(fsetup.bucket_storage),
 165                          BlockStorageTypeFactory(self._type_name))
 166         with open(fname, 'rb') as f:
 167             flen = len(f.read())
 168             self.assertEqual(
 169                 flen,
 170                 HeapStorage.compute_storage_size(
 171                     bsize,
 172                     heap_height,
 173                     blocks_per_bucket=blocks_per_bucket))
 174             self.assertEqual(
 175                 flen >
 176                 HeapStorage.compute_storage_size(
 177                     bsize,
 178                     heap_height,
 179                     blocks_per_bucket=blocks_per_bucket,
 180                     ignore_header=True),
 181                 True)
 182         with HeapStorage(
 183                 fname,
 184                 storage_type=self._type_name) as f:
 185             self.assertEqual(f.header_data, bytes())
 186             self.assertEqual(fsetup.header_data, bytes())
 187             self.assertEqual(f.blocks_per_bucket,
 188                              blocks_per_bucket)
 189             self.assertEqual(fsetup.blocks_per_bucket,
 190                              blocks_per_bucket)
 191             self.assertEqual(f.bucket_count,
 192                              2**(heap_height+1) - 1)
 193             self.assertEqual(fsetup.bucket_count,
 194                              2**(heap_height+1) - 1)
 195             self.assertEqual(f.bucket_size,
 196                              bsize * blocks_per_bucket)
 197             self.assertEqual(fsetup.bucket_size,
 198                              bsize * blocks_per_bucket)
 199             self.assertEqual(f.storage_name, fname)
 200             self.assertEqual(fsetup.storage_name, fname)
 201         os.remove(fname)
 202
 203     def test_setup_withdata(self):
 204         fname = ".".join(self.id().split(".")[1:])
 205         fname += ".bin"
 206         fname = os.path.join(thisdir, fname)
 207         if os.path.exists(fname):
 208             os.remove(fname)                           # pragma: no cover
 209         bsize = 10
 210         heap_height = 2
 211         blocks_per_bucket = 1
 212         header_data = bytes(bytearray([0,1,2]))
 213         fsetup = HeapStorage.setup(
 214             fname,
 215             bsize,
 216             heap_height,
 217             blocks_per_bucket=blocks_per_bucket,
 218             header_data=header_data)
 219         fsetup.close()
 220         self.assertEqual(type(fsetup.bucket_storage),
 221                          BlockStorageTypeFactory(self._type_name))
 222         with open(fname, 'rb') as f:
 223             flen = len(f.read())
 224             self.assertEqual(
 225                 flen,
 226                 HeapStorage.compute_storage_size(
 227                     bsize,
 228                     heap_height,
 229                     header_data=header_data))
 230             self.assertTrue(len(header_data) > 0)
 231             self.assertEqual(
 232                 HeapStorage.compute_storage_size(
 233                     bsize,
 234                     heap_height,
 235                     storage_type=self._type_name) <
 236                 HeapStorage.compute_storage_size(
 237                     bsize,
 238                     heap_height,
 239                     storage_type=self._type_name,
 240                     header_data=header_data),
 241                 True)
 242             self.assertEqual(
 243                 flen >
 244                 HeapStorage.compute_storage_size(
 245                     bsize,
 246                     heap_height,
 247                     storage_type=self._type_name,
 248                     header_data=header_data,
 249                     ignore_header=True),
 250                 True)
 251         with HeapStorage(
 252                 fname,
 253                 storage_type=self._type_name) as f:
 254             self.assertEqual(f.header_data, header_data)
 255             self.assertEqual(fsetup.header_data, header_data)
 256             self.assertEqual(f.blocks_per_bucket,
 257                              blocks_per_bucket)
 258             self.assertEqual(fsetup.blocks_per_bucket,
 259                              blocks_per_bucket)
 260             self.assertEqual(f.bucket_count,
 261                              2**(heap_height+1) - 1)
 262             self.assertEqual(fsetup.bucket_count,
 263                              2**(heap_height+1) - 1)
 264             self.assertEqual(f.bucket_size,
 265                              bsize * blocks_per_bucket)
 266             self.assertEqual(fsetup.bucket_size,
 267                              bsize * blocks_per_bucket)
 268             self.assertEqual(f.storage_name, fname)
 269             self.assertEqual(fsetup.storage_name, fname)
 270         os.remove(fname)
 271
 272     def test_init_noexists(self):
 273         self.assertEqual(os.path.exists(self._dummy_name), False)
 274         with self.assertRaises(IOError):
 275             with HeapStorage(
 276                     self._dummy_name,
 277                     storage_type=self._type_name) as f:
 278                 pass                                   # pragma: no cover
 279
 280     def test_init_exists(self):
 281         self.assertEqual(os.path.exists(self._testfname), True)
 282         with open(self._testfname, 'rb') as f:
 283             databefore = f.read()
 284         with self.assertRaises(ValueError):
 285             with BlockStorageFile(self._testfname) as fb:
 286                 with HeapStorage(fb, storage_type='file') as f:
 287                     pass                               # pragma: no cover
 288         with HeapStorage(
 289                 self._testfname,
 290                 storage_type=self._type_name) as f:
 291             self.assertEqual(f.bucket_size,
 292                              self._block_size * \
 293                              self._blocks_per_bucket)
 294             self.assertEqual(f.bucket_count,
 295                              self._bucket_count)
 296             self.assertEqual(f.storage_name, self._testfname)
 297             self.assertEqual(f.header_data, bytes())
 298         self.assertEqual(os.path.exists(self._testfname), True)
 299         with open(self._testfname, 'rb') as f:
 300             dataafter = f.read()
 301         self.assertEqual(databefore, dataafter)
 302
 303     def test_read_path(self):
 304
 305         with HeapStorage(
 306                 self._testfname,
 307                 storage_type=self._type_name) as f:
 308             self.assertEqual(f.bytes_sent, 0)
 309             self.assertEqual(f.bytes_received, 0)
 310
 311             self.assertEqual(
 312                 f.virtual_heap.first_bucket_at_level(0), 0)
 313             self.assertNotEqual(
 314                 f.virtual_heap.last_leaf_bucket(), 0)
 315             total_buckets = 0
 316             for b in range(f.virtual_heap.first_bucket_at_level(0),
 317                            f.virtual_heap.last_leaf_bucket()+1):
 318                 data = f.read_path(b)
 319                 bucket_path = f.virtual_heap.Node(b).\
 320                               bucket_path_from_root()
 321                 total_buckets += len(bucket_path)
 322                 self.assertEqual(f.virtual_heap.Node(b).level+1,
 323                                  len(bucket_path))
 324                 for i, bucket in zip(bucket_path, data):
 325                     self.assertEqual(list(bytearray(bucket)),
 326                                      list(self._buckets[i]))
 327
 328             self.assertEqual(f.bytes_sent, 0)
 329             self.assertEqual(f.bytes_received,
 330                              total_buckets*f.bucket_storage.block_size)
 331
 332     def test_write_path(self):
 333         data = [bytearray([self._bucket_count]) * \
 334                 self._block_size * \
 335                 self._blocks_per_bucket
 336                 for i in xrange(self._block_count)]
 337         with HeapStorage(
 338                 self._testfname,
 339                 storage_type=self._type_name) as f:
 340             self.assertEqual(f.bytes_sent, 0)
 341             self.assertEqual(f.bytes_received, 0)
 342
 343             self.assertEqual(
 344                 f.virtual_heap.first_bucket_at_level(0), 0)
 345             self.assertNotEqual(
 346                 f.virtual_heap.last_leaf_bucket(), 0)
 347             total_buckets = 0
 348             for b in range(f.virtual_heap.first_bucket_at_level(0),
 349                            f.virtual_heap.last_leaf_bucket()+1):
 350                 orig = f.read_path(b)
 351                 bucket_path = f.virtual_heap.Node(b).\
 352                               bucket_path_from_root()
 353                 total_buckets += len(bucket_path)
 354                 self.assertNotEqual(len(bucket_path), 0)
 355                 self.assertEqual(f.virtual_heap.Node(b).level+1,
 356                                  len(bucket_path))
 357                 self.assertEqual(len(orig), len(bucket_path))
 358                 for i, bucket in zip(bucket_path, orig):
 359                     self.assertEqual(list(bytearray(bucket)),
 360                                      list(self._buckets[i]))
 361                 f.write_path(b, [bytes(data[i])
 362                                  for i in bucket_path])
 363
 364                 new = f.read_path(b)
 365                 self.assertEqual(len(new), len(bucket_path))
 366                 for i, bucket in zip(bucket_path, new):
 367                     self.assertEqual(list(bytearray(bucket)),
 368                                      list(data[i]))
 369
 370                 f.write_path(b, [bytes(self._buckets[i])
 371                                  for i in bucket_path])
 372
 373                 orig = f.read_path(b)
 374                 self.assertEqual(len(orig), len(bucket_path))
 375                 for i, bucket in zip(bucket_path, orig):
 376                     self.assertEqual(list(bytearray(bucket)),
 377                                      list(self._buckets[i]))
 378
 379             self.assertEqual(f.bytes_sent,
 380                              total_buckets*f.bucket_storage.block_size*2)
 381             self.assertEqual(f.bytes_received,
 382                              total_buckets*f.bucket_storage.block_size*3)
 383
 384     def test_update_header_data(self):
 385         fname = ".".join(self.id().split(".")[1:])
 386         fname += ".bin"
 387         fname = os.path.join(thisdir, fname)
 388         if os.path.exists(fname):
 389             os.remove(fname)                           # pragma: no cover
 390         bsize = 10
 391         heap_height = 2
 392         blocks_per_bucket = 1
 393         header_data = bytes(bytearray([0,1,2]))
 394         fsetup = HeapStorage.setup(
 395             fname,
 396             block_size=bsize,
 397             heap_height=heap_height,
 398             blocks_per_bucket=blocks_per_bucket,
 399             header_data=header_data)
 400         fsetup.close()
 401         new_header_data = bytes(bytearray([1,1,1]))
 402         with HeapStorage(
 403                 fname,
 404                 storage_type=self._type_name) as f:
 405             self.assertEqual(f.header_data, header_data)
 406             f.update_header_data(new_header_data)
 407             self.assertEqual(f.header_data, new_header_data)
 408         with HeapStorage(
 409                 fname,
 410                 storage_type=self._type_name) as f:
 411             self.assertEqual(f.header_data, new_header_data)
 412         with self.assertRaises(ValueError):
 413             with HeapStorage(
 414                     fname,
 415                     storage_type=self._type_name) as f:
 416                 f.update_header_data(bytes(bytearray([1,1])))
 417         with self.assertRaises(ValueError):
 418             with HeapStorage(
 419                     fname,
 420                     storage_type=self._type_name) as f:
 421                 f.update_header_data(bytes(bytearray([1,1,1,1])))
 422         with HeapStorage(
 423                 fname,
 424                 storage_type=self._type_name) as f:
 425             self.assertEqual(f.header_data, new_header_data)
 426         os.remove(fname)
 427
 428     def test_locked_flag(self):
 429         with HeapStorage(self._testfname,
 430                                   storage_type=self._type_name) as f:
 431             with self.assertRaises(IOError):
 432                 with HeapStorage(self._testfname,
 433                                           storage_type=self._type_name) as f1:
 434                     pass                               # pragma: no cover
 435             with self.assertRaises(IOError):
 436                 with HeapStorage(self._testfname,
 437                                           storage_type=self._type_name) as f1:
 438                     pass                               # pragma: no cover
 439             with HeapStorage(self._testfname,
 440                                       storage_type=self._type_name,
 441                                       ignore_lock=True) as f1:
 442                 pass
 443             with self.assertRaises(IOError):
 444                 with HeapStorage(self._testfname,
 445                                           storage_type=self._type_name) as f1:
 446                     pass                               # pragma: no cover
 447             with HeapStorage(self._testfname,
 448                                       storage_type=self._type_name,
 449                                       ignore_lock=True) as f1:
 450                 pass
 451             with HeapStorage(self._testfname,
 452                                       storage_type=self._type_name,
 453                                       ignore_lock=True) as f1:
 454                 pass
 455         with HeapStorage(self._testfname,
 456                                   storage_type=self._type_name) as f:
 457             pass
 458
 459     def test_read_path_cloned(self):
 460
 461         with HeapStorage(
 462                 self._testfname,
 463                 storage_type=self._type_name) as forig:
 464             self.assertEqual(forig.bytes_sent, 0)
 465             self.assertEqual(forig.bytes_received, 0)
 466             with forig.clone_device() as f:
 467                 self.assertEqual(forig.bytes_sent, 0)
 468                 self.assertEqual(forig.bytes_received, 0)
 469                 self.assertEqual(f.bytes_sent, 0)
 470                 self.assertEqual(f.bytes_received, 0)
 471
 472                 self.assertEqual(
 473                     f.virtual_heap.first_bucket_at_level(0), 0)
 474                 self.assertNotEqual(
 475                     f.virtual_heap.last_leaf_bucket(), 0)
 476                 total_buckets = 0
 477                 for b in range(f.virtual_heap.first_bucket_at_level(0),
 478                                f.virtual_heap.last_leaf_bucket()+1):
 479                     data = f.read_path(b)
 480                     bucket_path = f.virtual_heap.Node(b).\
 481                                   bucket_path_from_root()
 482                     total_buckets += len(bucket_path)
 483                     self.assertEqual(f.virtual_heap.Node(b).level+1,
 484                                      len(bucket_path))
 485                     for i, bucket in zip(bucket_path, data):
 486                         self.assertEqual(list(bytearray(bucket)),
 487                                          list(self._buckets[i]))
 488
 489                 self.assertEqual(f.bytes_sent, 0)
 490                 self.assertEqual(f.bytes_received,
 491                                  total_buckets*f.bucket_storage.block_size)
 492             self.assertEqual(forig.bytes_sent, 0)
 493             self.assertEqual(forig.bytes_received, 0)
 494
 495     def test_write_path_cloned(self):
 496         data = [bytearray([self._bucket_count]) * \
 497                 self._block_size * \
 498                 self._blocks_per_bucket
 499                 for i in xrange(self._block_count)]
 500         with HeapStorage(
 501                 self._testfname,
 502                 storage_type=self._type_name) as forig:
 503             self.assertEqual(forig.bytes_sent, 0)
 504             self.assertEqual(forig.bytes_received, 0)
 505             with forig.clone_device() as f:
 506                 self.assertEqual(forig.bytes_sent, 0)
 507                 self.assertEqual(forig.bytes_received, 0)
 508                 self.assertEqual(f.bytes_sent, 0)
 509                 self.assertEqual(f.bytes_received, 0)
 510
 511                 self.assertEqual(
 512                     f.virtual_heap.first_bucket_at_level(0), 0)
 513                 self.assertNotEqual(
 514                     f.virtual_heap.last_leaf_bucket(), 0)
 515                 total_buckets = 0
 516                 for b in range(f.virtual_heap.first_bucket_at_level(0),
 517                                f.virtual_heap.last_leaf_bucket()+1):
 518                     orig = f.read_path(b)
 519                     bucket_path = f.virtual_heap.Node(b).\
 520                                   bucket_path_from_root()
 521                     total_buckets += len(bucket_path)
 522                     self.assertNotEqual(len(bucket_path), 0)
 523                     self.assertEqual(f.virtual_heap.Node(b).level+1,
 524                                      len(bucket_path))
 525                     self.assertEqual(len(orig), len(bucket_path))
 526                     for i, bucket in zip(bucket_path, orig):
 527                         self.assertEqual(list(bytearray(bucket)),
 528                                          list(self._buckets[i]))
 529                     f.write_path(b, [bytes(data[i])
 530                                      for i in bucket_path])
 531
 532                     new = f.read_path(b)
 533                     self.assertEqual(len(new), len(bucket_path))
 534                     for i, bucket in zip(bucket_path, new):
 535                         self.assertEqual(list(bytearray(bucket)),
 536                                          list(data[i]))
 537
 538                     f.write_path(b, [bytes(self._buckets[i])
 539                                      for i in bucket_path])
 540
 541                     orig = f.read_path(b)
 542                     self.assertEqual(len(orig), len(bucket_path))
 543                     for i, bucket in zip(bucket_path, orig):
 544                         self.assertEqual(list(bytearray(bucket)),
 545                                          list(self._buckets[i]))
 546
 547                 self.assertEqual(f.bytes_sent,
 548                                  total_buckets*f.bucket_storage.block_size*2)
 549                 self.assertEqual(f.bytes_received,
 550                                  total_buckets*f.bucket_storage.block_size*3)
 551             self.assertEqual(forig.bytes_sent, 0)
 552             self.assertEqual(forig.bytes_received, 0)
 553
 554 if __name__ == "__main__":
 555     unittest2.main()                                    # pragma: no cover