PyORAM/src/pyoram/storage/block_storage_ram.py

   1 __all__ = ('BlockStorageRAM',)
   2
   3 import os
   4 import struct
   5 import logging
   6 import errno
   7 from multiprocessing.pool import ThreadPool
   8
   9 import pyoram
  10 from pyoram.storage.block_storage import \
  11     (BlockStorageInterface,
  12      BlockStorageTypeFactory)
  13 from pyoram.storage.block_storage_mmap import \
  14     (BlockStorageMMap,
  15      _BlockStorageMemoryImpl)
  16
  17 import tqdm
  18 import six
  19 from six.moves import xrange
  20
  21 log = logging.getLogger("pyoram")
  22
  23 class BlockStorageRAM(_BlockStorageMemoryImpl,
  24                       BlockStorageInterface):
  25     """
  26     A class implementing the block storage interface where all data is
  27     kept in RAM. This class uses the same storage format as
  28     BlockStorageFile. Thus, a block storage space can be created using
  29     this class and then, after saving the raw storage data to disk,
  30     reopened with any other class compatible with BlockStorageFile
  31     (and visa versa).
  32     """
  33
  34     _index_struct_string = BlockStorageMMap._index_struct_string
  35     _index_offset = struct.calcsize(_index_struct_string)
  36
  37     def __init__(self,
  38                  storage_data,
  39                  threadpool_size=None,
  40                  ignore_lock=False):
  41
  42         self._bytes_sent = 0
  43         self._bytes_received = 0
  44         self._ignore_lock = ignore_lock
  45         self._f = None
  46         self._pool = None
  47         self._close_pool = True
  48         if type(storage_data) is not bytearray:
  49             raise TypeError(
  50                 "BlockStorageRAM requires input argument of type "
  51                 "'bytearray'. Invalid input type: %s"
  52                 % (type(storage_data)))
  53         self._f = storage_data
  54         self._block_size, self._block_count, user_header_size, locked = \
  55             struct.unpack(
  56                 BlockStorageRAM._index_struct_string,
  57                 self._f[:BlockStorageRAM._index_offset])
  58
  59         if locked and (not self._ignore_lock):
  60             raise IOError(
  61                 "Can not open block storage device because it is "
  62                 "locked by another process. To ignore this check, "
  63                 "initialize this class with the keyword 'ignore_lock' "
  64                 "set to True.")
  65         self._user_header_data = bytes()
  66         if user_header_size > 0:
  67             self._user_header_data = \
  68                 bytes(self._f[BlockStorageRAM._index_offset:\
  69                               (BlockStorageRAM._index_offset+user_header_size)])
  70         assert len(self._user_header_data) == user_header_size
  71         self._header_offset = BlockStorageRAM._index_offset + \
  72                               len(self._user_header_data)
  73
  74         if not self._ignore_lock:
  75             # turn on the locked flag
  76             self._f[:BlockStorageRAM._index_offset] = \
  77                 struct.pack(BlockStorageRAM._index_struct_string,
  78                             self.block_size,
  79                             self.block_count,
  80                             len(self._user_header_data),
  81                             True)
  82
  83         # Although we do not use the threadpool we still
  84         # create just in case we are the first
  85         if threadpool_size != 0:
  86             self._pool = ThreadPool(threadpool_size)
  87
  88     #
  89     # Add some methods specific to BlockStorageRAM
  90     #
  91
  92     @staticmethod
  93     def fromfile(file_,
  94                  threadpool_size=None,
  95                  ignore_lock=False):
  96         """
  97         Instantiate BlockStorageRAM device from a file saved in block
  98         storage format. The file_ argument can be a file object or a
  99         string that represents a filename. If called with a file
 100         object, it should be opened in binary mode, and the caller is
 101         responsible for closing the file.
 102
 103         This method returns a BlockStorageRAM instance.
 104         """
 105         close_file = False
 106         if not hasattr(file_, 'read'):
 107             file_ = open(file_, 'rb')
 108             close_file = True
 109         try:
 110             header_data = file_.read(BlockStorageRAM._index_offset)
 111             block_size, block_count, user_header_size, locked = \
 112                 struct.unpack(
 113                     BlockStorageRAM._index_struct_string,
 114                     header_data)
 115             if locked and (not ignore_lock):
 116                 raise IOError(
 117                     "Can not open block storage device because it is "
 118                     "locked by another process. To ignore this check, "
 119                     "call this method with the keyword 'ignore_lock' "
 120                     "set to True.")
 121             header_offset = len(header_data) + \
 122                             user_header_size
 123             f = bytearray(header_offset + \
 124                           (block_size * block_count))
 125             f[:header_offset] = header_data + file_.read(user_header_size)
 126             f[header_offset:] = file_.read(block_size * block_count)
 127         finally:
 128             if close_file:
 129                 file_.close()
 130
 131         return BlockStorageRAM(f,
 132                                threadpool_size=threadpool_size,
 133                                ignore_lock=ignore_lock)
 134
 135     def tofile(self, file_):
 136         """
 137         Dump all storage data to a file. The file_ argument can be a
 138         file object or a string that represents a filename. If called
 139         with a file object, it should be opened in binary mode, and
 140         the caller is responsible for closing the file.
 141
 142         The method should only be called after the storage device has
 143         been closed to ensure that the locked flag has been set to
 144         False.
 145         """
 146         close_file = False
 147         if not hasattr(file_, 'write'):
 148             file_ = open(file_, 'wb')
 149             close_file = True
 150         file_.write(self._f)
 151         if close_file:
 152             file_.close()
 153
 154     @property
 155     def data(self):
 156         """Access the raw bytearray"""
 157         return self._f
 158
 159     #
 160     # Define BlockStorageInterface Methods
 161     #
 162
 163     def clone_device(self):
 164         f = BlockStorageRAM(self._f,
 165                             threadpool_size=0,
 166                             ignore_lock=True)
 167         f._pool = self._pool
 168         f._close_pool = False
 169         return f
 170
 171     @classmethod
 172     def compute_storage_size(cls, *args, **kwds):
 173         return BlockStorageMMap.compute_storage_size(*args, **kwds)
 174
 175     @classmethod
 176     def setup(cls,
 177               storage_name,
 178               block_size,
 179               block_count,
 180               initialize=None,
 181               header_data=None,
 182               ignore_existing=False,
 183               threadpool_size=None):
 184
 185         # We ignore the 'storage_name' argument
 186         # We ignore the 'ignore_existing' flag
 187         if (block_size <= 0) or (block_size != int(block_size)):
 188             raise ValueError(
 189                 "Block size (bytes) must be a positive integer: %s"
 190                 % (block_size))
 191         if (block_count <= 0) or (block_count != int(block_count)):
 192             raise ValueError(
 193                 "Block count must be a positive integer: %s"
 194                 % (block_count))
 195         if (header_data is not None) and \
 196            (type(header_data) is not bytes):
 197             raise TypeError(
 198                 "'header_data' must be of type bytes. "
 199                 "Invalid type: %s" % (type(header_data)))
 200
 201         if initialize is None:
 202             zeros = bytes(bytearray(block_size))
 203             initialize = lambda i: zeros
 204
 205         # create_index
 206         index_data = None
 207         if header_data is None:
 208             index_data = struct.pack(BlockStorageRAM._index_struct_string,
 209                                      block_size,
 210                                      block_count,
 211                                      0,
 212                                      False)
 213             header_data = bytes()
 214         else:
 215             index_data = struct.pack(BlockStorageRAM._index_struct_string,
 216                                      block_size,
 217                                      block_count,
 218                                      len(header_data),
 219                                      False)
 220         header_offset = len(index_data) + len(header_data)
 221         f = bytearray(header_offset + \
 222                       (block_size * block_count))
 223         f[:header_offset] = index_data + header_data
 224         progress_bar = tqdm.tqdm(total=block_count*block_size,
 225                                  desc="Initializing File Block Storage Space",
 226                                  unit="B",
 227                                  unit_scale=True,
 228                                  disable=not pyoram.config.SHOW_PROGRESS_BAR)
 229         for i in xrange(block_count):
 230             block = initialize(i)
 231             assert len(block) == block_size, \
 232                 ("%s != %s" % (len(block), block_size))
 233             pos_start = header_offset + i * block_size
 234             pos_start = header_offset + i * block_size
 235             pos_stop = pos_start + block_size
 236             f[pos_start:pos_stop] = block[:]
 237             progress_bar.update(n=block_size)
 238         progress_bar.close()
 239
 240         return BlockStorageRAM(f, threadpool_size=threadpool_size)
 241
 242     @property
 243     def header_data(self):
 244         return self._user_header_data
 245
 246     @property
 247     def block_count(self):
 248         return self._block_count
 249
 250     @property
 251     def block_size(self):
 252         return self._block_size
 253
 254     @property
 255     def storage_name(self):
 256         return None
 257
 258     def update_header_data(self, new_header_data):
 259         if len(new_header_data) != len(self.header_data):
 260             raise ValueError(
 261                 "The size of header data can not change.\n"
 262                 "Original bytes: %s\n"
 263                 "New bytes: %s" % (len(self.header_data),
 264                                    len(new_header_data)))
 265         self._user_header_data = bytes(new_header_data)
 266         self._f[BlockStorageRAM._index_offset:\
 267                 (BlockStorageRAM._index_offset+len(new_header_data))] = \
 268             self._user_header_data
 269
 270     def close(self):
 271         if self._close_pool and (self._pool is not None):
 272             self._pool.close()
 273             self._pool.join()
 274             self._pool = None
 275         if not self._ignore_lock:
 276             # turn off the locked flag
 277             self._f[:BlockStorageRAM._index_offset] = \
 278                 struct.pack(BlockStorageRAM._index_struct_string,
 279                             self.block_size,
 280                             self.block_count,
 281                             len(self._user_header_data),
 282                             False)
 283             self._ignore_lock = True
 284
 285     #
 286     # We must cast from bytearray to bytes
 287     # when reading from a bytearray so that this
 288     # class works with the encryption layer.
 289     #
 290
 291     def read_blocks(self, indices):
 292         return [bytes(block) for block
 293                 in super(BlockStorageRAM, self).read_blocks(indices)]
 294
 295     def yield_blocks(self, indices):
 296         for block in super(BlockStorageRAM, self).yield_blocks(indices):
 297             yield bytes(block)
 298
 299     def read_block(self, i):
 300         return bytes(super(BlockStorageRAM, self).read_block(i))
 301
 302     #def write_blocks(...)
 303
 304     #def write_block(...)
 305
 306     @property
 307     def bytes_sent(self):
 308         return self._bytes_sent
 309
 310     @property
 311     def bytes_received(self):
 312         return self._bytes_received
 313
 314 BlockStorageTypeFactory.register_device("ram", BlockStorageRAM)