1 __all__ = ('BlockStorageRAM',)
7 from multiprocessing.pool import ThreadPool
10 from pyoram.storage.block_storage import \
11 (BlockStorageInterface,
12 BlockStorageTypeFactory)
13 from pyoram.storage.block_storage_mmap import \
15 _BlockStorageMemoryImpl)
19 from six.moves import xrange
21 log = logging.getLogger("pyoram")
23 class BlockStorageRAM(_BlockStorageMemoryImpl,
24 BlockStorageInterface):
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
34 _index_struct_string = BlockStorageMMap._index_struct_string
35 _index_offset = struct.calcsize(_index_struct_string)
43 self._bytes_received = 0
44 self._ignore_lock = ignore_lock
47 self._close_pool = True
48 if type(storage_data) is not bytearray:
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 = \
56 BlockStorageRAM._index_struct_string,
57 self._f[:BlockStorageRAM._index_offset])
59 if locked and (not self._ignore_lock):
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' "
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)
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,
80 len(self._user_header_data),
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)
89 # Add some methods specific to BlockStorageRAM
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.
103 This method returns a BlockStorageRAM instance.
106 if not hasattr(file_, 'read'):
107 file_ = open(file_, 'rb')
110 header_data = file_.read(BlockStorageRAM._index_offset)
111 block_size, block_count, user_header_size, locked = \
113 BlockStorageRAM._index_struct_string,
115 if locked and (not ignore_lock):
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' "
121 header_offset = len(header_data) + \
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)
131 return BlockStorageRAM(f,
132 threadpool_size=threadpool_size,
133 ignore_lock=ignore_lock)
135 def tofile(self, file_):
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.
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
147 if not hasattr(file_, 'write'):
148 file_ = open(file_, 'wb')
156 """Access the raw bytearray"""
160 # Define BlockStorageInterface Methods
163 def clone_device(self):
164 f = BlockStorageRAM(self._f,
168 f._close_pool = False
172 def compute_storage_size(cls, *args, **kwds):
173 return BlockStorageMMap.compute_storage_size(*args, **kwds)
182 ignore_existing=False,
183 threadpool_size=None):
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)):
189 "Block size (bytes) must be a positive integer: %s"
191 if (block_count <= 0) or (block_count != int(block_count)):
193 "Block count must be a positive integer: %s"
195 if (header_data is not None) and \
196 (type(header_data) is not bytes):
198 "'header_data' must be of type bytes. "
199 "Invalid type: %s" % (type(header_data)))
201 if initialize is None:
202 zeros = bytes(bytearray(block_size))
203 initialize = lambda i: zeros
207 if header_data is None:
208 index_data = struct.pack(BlockStorageRAM._index_struct_string,
213 header_data = bytes()
215 index_data = struct.pack(BlockStorageRAM._index_struct_string,
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",
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)
240 return BlockStorageRAM(f, threadpool_size=threadpool_size)
243 def header_data(self):
244 return self._user_header_data
247 def block_count(self):
248 return self._block_count
251 def block_size(self):
252 return self._block_size
255 def storage_name(self):
258 def update_header_data(self, new_header_data):
259 if len(new_header_data) != len(self.header_data):
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
271 if self._close_pool and (self._pool is not 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,
281 len(self._user_header_data),
283 self._ignore_lock = True
286 # We must cast from bytearray to bytes
287 # when reading from a bytearray so that this
288 # class works with the encryption layer.
291 def read_blocks(self, indices):
292 return [bytes(block) for block
293 in super(BlockStorageRAM, self).read_blocks(indices)]
295 def yield_blocks(self, indices):
296 for block in super(BlockStorageRAM, self).yield_blocks(indices):
299 def read_block(self, i):
300 return bytes(super(BlockStorageRAM, self).read_block(i))
302 #def write_blocks(...)
304 #def write_block(...)
307 def bytes_sent(self):
308 return self._bytes_sent
311 def bytes_received(self):
312 return self._bytes_received
314 BlockStorageTypeFactory.register_device("ram", BlockStorageRAM)