Viewing file:  images.py (18.16 KB) -rw-r--r--
Select action/file-type:
 (+) |  (+) |  (+) | Code (+) | Session (+) |  (+) | SDB (+) |  (+) |  (+) |  (+) |  (+) |  (+) |
# This file contains methods to deal with criu images.
#
# According to http://criu.org/Images, criu images can be described
# with such IOW:
#
# IMAGE_FILE ::= MAGIC { ENTRY }
# ENTRY ::= SIZE PAYLOAD [ EXTRA ]
# PAYLOAD ::= "message encoded in ProtocolBuffer format"
# EXTRA ::= "arbitrary blob, depends on the PAYLOAD contents"
#
# MAGIC ::= "32 bit integer"
# SIZE ::= "32 bit integer, equals the PAYLOAD length"
#
# Images v1.1 NOTE: MAGIC now consist of 2 32 bit integers, first one is
# MAGIC_COMMON or MAGIC_SERVICE and the second one is same as MAGIC
# in images V1.0. We don't keep "first" magic in json images.
#
# In order to convert images to human-readable format, we use dict(json).
# Using json not only allows us to easily read\write images, but also
# to use a great variety of tools out there to manipulate them.
# It also allows us to clearly describe criu images structure.
#
# Using dict(json) format, criu images can be described like:
#
# {
# 'magic' : 'FOO',
# 'entries' : [
# entry,
# ...
# ]
# }
#
# Entry, in its turn, could be described as:
#
# {
# pb_msg,
# 'extra' : extra_msg
# }
#
import io
import base64
import struct
import os
import array
from . import magic
from . import pb
from . import pb2dict
if "encodebytes" not in dir(base64):
base64.encodebytes = base64.encodestring
base64.decodebytes = base64.decodestring
#
# Predefined hardcoded constants
sizeof_u16 = 2
sizeof_u32 = 4
sizeof_u64 = 8
# A helper for rounding
def round_up(x, y):
return (((x - 1) | (y - 1)) + 1)
class MagicException(Exception):
def __init__(self, magic):
self.magic = magic
# Generic class to handle loading/dumping criu images entries from/to bin
# format to/from dict(json).
class entry_handler:
"""
Generic class to handle loading/dumping criu images
entries from/to bin format to/from dict(json).
"""
def __init__(self, payload, extra_handler=None):
"""
Sets payload class and extra handler class.
"""
self.payload = payload
self.extra_handler = extra_handler
def load(self, f, pretty=False, no_payload=False):
"""
Convert criu image entries from binary format to dict(json).
Takes a file-like object and returnes a list with entries in
dict(json) format.
"""
entries = []
while True:
entry = {}
# Read payload
pbuff = self.payload()
buf = f.read(4)
if buf == b'':
break
size, = struct.unpack('i', buf)
pbuff.ParseFromString(f.read(size))
entry = pb2dict.pb2dict(pbuff, pretty)
# Read extra
if self.extra_handler:
if no_payload:
def human_readable(num):
for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
if num < 1024.0:
if int(num) == num:
return "%d%sB" % (num, unit)
else:
return "%.1f%sB" % (num, unit)
num /= 1024.0
return "%.1fYB" % num
pl_size = self.extra_handler.skip(f, pbuff)
entry['extra'] = '... <%s>' % human_readable(pl_size)
else:
entry['extra'] = self.extra_handler.load(f, pbuff)
entries.append(entry)
return entries
def loads(self, s, pretty=False):
"""
Same as load(), but takes a string as an argument.
"""
f = io.BytesIO(s)
return self.load(f, pretty)
def dump(self, entries, f):
"""
Convert criu image entries from dict(json) format to binary.
Takes a list of entries and a file-like object to write entries
in binary format to.
"""
for entry in entries:
extra = entry.pop('extra', None)
# Write payload
pbuff = self.payload()
pb2dict.dict2pb(entry, pbuff)
pb_str = pbuff.SerializeToString()
size = len(pb_str)
f.write(struct.pack('i', size))
f.write(pb_str)
# Write extra
if self.extra_handler and extra:
self.extra_handler.dump(extra, f, pbuff)
def dumps(self, entries):
"""
Same as dump(), but doesn't take file-like object and just
returns a string.
"""
f = io.BytesIO('')
self.dump(entries, f)
return f.read()
def count(self, f):
"""
Counts the number of top-level object in the image file
"""
entries = 0
while True:
buf = f.read(4)
if buf == '':
break
size, = struct.unpack('i', buf)
f.seek(size, 1)
entries += 1
return entries
# Special handler for pagemap.img
class pagemap_handler:
"""
Special entry handler for pagemap.img, which is unique in a way
that it has a header of pagemap_head type followed by entries
of pagemap_entry type.
"""
def load(self, f, pretty=False, no_payload=False):
entries = []
pbuff = pb.pagemap_head()
while True:
buf = f.read(4)
if buf == b'':
break
size, = struct.unpack('i', buf)
pbuff.ParseFromString(f.read(size))
entries.append(pb2dict.pb2dict(pbuff, pretty))
pbuff = pb.pagemap_entry()
return entries
def loads(self, s, pretty=False):
f = io.BytesIO(s)
return self.load(f, pretty)
def dump(self, entries, f):
pbuff = pb.pagemap_head()
for item in entries:
pb2dict.dict2pb(item, pbuff)
pb_str = pbuff.SerializeToString()
size = len(pb_str)
f.write(struct.pack('i', size))
f.write(pb_str)
pbuff = pb.pagemap_entry()
def dumps(self, entries):
f = io.BytesIO('')
self.dump(entries, f)
return f.read()
def count(self, f):
return entry_handler(None).count(f) - 1
# Special handler for ghost-file.img
class ghost_file_handler:
def load(self, f, pretty=False, no_payload=False):
entries = []
gf = pb.ghost_file_entry()
buf = f.read(4)
size, = struct.unpack('i', buf)
gf.ParseFromString(f.read(size))
g_entry = pb2dict.pb2dict(gf, pretty)
if gf.chunks:
entries.append(g_entry)
while True:
gc = pb.ghost_chunk_entry()
buf = f.read(4)
if buf == '':
break
size, = struct.unpack('i', buf)
gc.ParseFromString(f.read(size))
entry = pb2dict.pb2dict(gc, pretty)
if no_payload:
f.seek(gc.len, os.SEEK_CUR)
else:
entry['extra'] = base64.encodebytes(f.read(gc.len))
entries.append(entry)
else:
if no_payload:
f.seek(0, os.SEEK_END)
else:
g_entry['extra'] = base64.encodebytes(f.read())
entries.append(g_entry)
return entries
def loads(self, s, pretty=False):
f = io.BytesIO(s)
return self.load(f, pretty)
def dump(self, entries, f):
pbuff = pb.ghost_file_entry()
item = entries.pop(0)
pb2dict.dict2pb(item, pbuff)
pb_str = pbuff.SerializeToString()
size = len(pb_str)
f.write(struct.pack('i', size))
f.write(pb_str)
if pbuff.chunks:
for item in entries:
pbuff = pb.ghost_chunk_entry()
pb2dict.dict2pb(item, pbuff)
pb_str = pbuff.SerializeToString()
size = len(pb_str)
f.write(struct.pack('i', size))
f.write(pb_str)
f.write(base64.decodebytes(item['extra']))
else:
f.write(base64.decodebytes(item['extra']))
def dumps(self, entries):
f = io.BytesIO('')
self.dump(entries, f)
return f.read()
# In following extra handlers we use base64 encoding
# to store binary data. Even though, the nature
# of base64 is that it increases the total size,
# it doesn't really matter, because our images
# do not store big amounts of binary data. They
# are negligible comparing to pages size.
class pipes_data_extra_handler:
def load(self, f, pload):
size = pload.bytes
data = f.read(size)
return base64.encodebytes(data)
def dump(self, extra, f, pload):
data = base64.decodebytes(extra)
f.write(data)
def skip(self, f, pload):
f.seek(pload.bytes, os.SEEK_CUR)
return pload.bytes
class sk_queues_extra_handler:
def load(self, f, pload):
size = pload.length
data = f.read(size)
return base64.encodebytes(data)
def dump(self, extra, f, _unused):
data = base64.decodebytes(extra)
f.write(data)
def skip(self, f, pload):
f.seek(pload.length, os.SEEK_CUR)
return pload.length
class tcp_stream_extra_handler:
def load(self, f, pbuff):
d = {}
inq = f.read(pbuff.inq_len)
outq = f.read(pbuff.outq_len)
d['inq'] = base64.encodebytes(inq)
d['outq'] = base64.encodebytes(outq)
return d
def dump(self, extra, f, _unused):
inq = base64.decodebytes(extra['inq'])
outq = base64.decodebytes(extra['outq'])
f.write(inq)
f.write(outq)
def skip(self, f, pbuff):
f.seek(0, os.SEEK_END)
return pbuff.inq_len + pbuff.outq_len
class ipc_sem_set_handler:
def load(self, f, pbuff):
entry = pb2dict.pb2dict(pbuff)
size = sizeof_u16 * entry['nsems']
rounded = round_up(size, sizeof_u64)
s = array.array('H')
if s.itemsize != sizeof_u16:
raise Exception("Array size mismatch")
s.fromstring(f.read(size))
f.seek(rounded - size, 1)
return s.tolist()
def dump(self, extra, f, pbuff):
entry = pb2dict.pb2dict(pbuff)
size = sizeof_u16 * entry['nsems']
rounded = round_up(size, sizeof_u64)
s = array.array('H')
if s.itemsize != sizeof_u16:
raise Exception("Array size mismatch")
s.fromlist(extra)
if len(s) != entry['nsems']:
raise Exception("Number of semaphores mismatch")
f.write(s.tostring())
f.write('\0' * (rounded - size))
def skip(self, f, pbuff):
entry = pb2dict.pb2dict(pbuff)
size = sizeof_u16 * entry['nsems']
f.seek(round_up(size, sizeof_u64), os.SEEK_CUR)
return size
class ipc_msg_queue_handler:
def load(self, f, pbuff):
entry = pb2dict.pb2dict(pbuff)
messages = []
for x in range(0, entry['qnum']):
buf = f.read(4)
if buf == '':
break
size, = struct.unpack('i', buf)
msg = pb.ipc_msg()
msg.ParseFromString(f.read(size))
rounded = round_up(msg.msize, sizeof_u64)
data = f.read(msg.msize)
f.seek(rounded - msg.msize, 1)
messages.append(pb2dict.pb2dict(msg))
messages.append(base64.encodebytes(data))
return messages
def dump(self, extra, f, pbuff):
entry = pb2dict.pb2dict(pbuff)
for i in range(0, len(extra), 2):
msg = pb.ipc_msg()
pb2dict.dict2pb(extra[i], msg)
msg_str = msg.SerializeToString()
size = len(msg_str)
f.write(struct.pack('i', size))
f.write(msg_str)
rounded = round_up(msg.msize, sizeof_u64)
data = base64.decodebytes(extra[i + 1])
f.write(data[:msg.msize])
f.write('\0' * (rounded - msg.msize))
def skip(self, f, pbuff):
entry = pb2dict.pb2dict(pbuff)
pl_len = 0
for x in range(0, entry['qnum']):
buf = f.read(4)
if buf == '':
break
size, = struct.unpack('i', buf)
msg = pb.ipc_msg()
msg.ParseFromString(f.read(size))
rounded = round_up(msg.msize, sizeof_u64)
f.seek(rounded, os.SEEK_CUR)
pl_len += size + msg.msize
return pl_len
class ipc_shm_handler:
def load(self, f, pbuff):
entry = pb2dict.pb2dict(pbuff)
size = entry['size']
data = f.read(size)
rounded = round_up(size, sizeof_u32)
f.seek(rounded - size, 1)
return base64.encodebytes(data)
def dump(self, extra, f, pbuff):
entry = pb2dict.pb2dict(pbuff)
size = entry['size']
data = base64.decodebytes(extra)
rounded = round_up(size, sizeof_u32)
f.write(data[:size])
f.write('\0' * (rounded - size))
def skip(self, f, pbuff):
entry = pb2dict.pb2dict(pbuff)
size = entry['size']
rounded = round_up(size, sizeof_u32)
f.seek(rounded, os.SEEK_CUR)
return size
handlers = {
'INVENTORY': entry_handler(pb.inventory_entry),
'CORE': entry_handler(pb.core_entry),
'IDS': entry_handler(pb.task_kobj_ids_entry),
'CREDS': entry_handler(pb.creds_entry),
'UTSNS': entry_handler(pb.utsns_entry),
'IPC_VAR': entry_handler(pb.ipc_var_entry),
'FS': entry_handler(pb.fs_entry),
'GHOST_FILE': ghost_file_handler(),
'MM': entry_handler(pb.mm_entry),
'CGROUP': entry_handler(pb.cgroup_entry),
'TCP_STREAM': entry_handler(pb.tcp_stream_entry,
tcp_stream_extra_handler()),
'STATS': entry_handler(pb.stats_entry),
'PAGEMAP': pagemap_handler(), # Special one
'PSTREE': entry_handler(pb.pstree_entry),
'REG_FILES': entry_handler(pb.reg_file_entry),
'NS_FILES': entry_handler(pb.ns_file_entry),
'EVENTFD_FILE': entry_handler(pb.eventfd_file_entry),
'EVENTPOLL_FILE': entry_handler(pb.eventpoll_file_entry),
'EVENTPOLL_TFD': entry_handler(pb.eventpoll_tfd_entry),
'SIGNALFD': entry_handler(pb.signalfd_entry),
'TIMERFD': entry_handler(pb.timerfd_entry),
'INOTIFY_FILE': entry_handler(pb.inotify_file_entry),
'INOTIFY_WD': entry_handler(pb.inotify_wd_entry),
'FANOTIFY_FILE': entry_handler(pb.fanotify_file_entry),
'FANOTIFY_MARK': entry_handler(pb.fanotify_mark_entry),
'VMAS': entry_handler(pb.vma_entry),
'PIPES': entry_handler(pb.pipe_entry),
'FIFO': entry_handler(pb.fifo_entry),
'SIGACT': entry_handler(pb.sa_entry),
'NETLINK_SK': entry_handler(pb.netlink_sk_entry),
'REMAP_FPATH': entry_handler(pb.remap_file_path_entry),
'MNTS': entry_handler(pb.mnt_entry),
'TTY_FILES': entry_handler(pb.tty_file_entry),
'TTY_INFO': entry_handler(pb.tty_info_entry),
'TTY_DATA': entry_handler(pb.tty_data_entry),
'RLIMIT': entry_handler(pb.rlimit_entry),
'TUNFILE': entry_handler(pb.tunfile_entry),
'EXT_FILES': entry_handler(pb.ext_file_entry),
'IRMAP_CACHE': entry_handler(pb.irmap_cache_entry),
'FILE_LOCKS': entry_handler(pb.file_lock_entry),
'FDINFO': entry_handler(pb.fdinfo_entry),
'UNIXSK': entry_handler(pb.unix_sk_entry),
'INETSK': entry_handler(pb.inet_sk_entry),
'PACKETSK': entry_handler(pb.packet_sock_entry),
'ITIMERS': entry_handler(pb.itimer_entry),
'POSIX_TIMERS': entry_handler(pb.posix_timer_entry),
'NETDEV': entry_handler(pb.net_device_entry),
'PIPES_DATA': entry_handler(pb.pipe_data_entry,
pipes_data_extra_handler()),
'FIFO_DATA': entry_handler(pb.pipe_data_entry, pipes_data_extra_handler()),
'SK_QUEUES': entry_handler(pb.sk_packet_entry, sk_queues_extra_handler()),
'IPCNS_SHM': entry_handler(pb.ipc_shm_entry, ipc_shm_handler()),
'IPCNS_SEM': entry_handler(pb.ipc_sem_entry, ipc_sem_set_handler()),
'IPCNS_MSG': entry_handler(pb.ipc_msg_entry, ipc_msg_queue_handler()),
'NETNS': entry_handler(pb.netns_entry),
'USERNS': entry_handler(pb.userns_entry),
'SECCOMP': entry_handler(pb.seccomp_entry),
'AUTOFS': entry_handler(pb.autofs_entry),
'FILES': entry_handler(pb.file_entry),
'CPUINFO': entry_handler(pb.cpuinfo_entry),
}
def __rhandler(f):
# Images v1.1 NOTE: First read "first" magic.
img_magic, = struct.unpack('i', f.read(4))
if img_magic in (magic.by_name['IMG_COMMON'],
magic.by_name['IMG_SERVICE']):
img_magic, = struct.unpack('i', f.read(4))
try:
m = magic.by_val[img_magic]
except:
raise MagicException(img_magic)
try:
handler = handlers[m]
except:
raise Exception("No handler found for image with magic " + m)
return m, handler
def load(f, pretty=False, no_payload=False):
"""
Convert criu image from binary format to dict(json).
Takes a file-like object to read criu image from.
Returns criu image in dict(json) format.
"""
image = {}
m, handler = __rhandler(f)
image['magic'] = m
image['entries'] = handler.load(f, pretty, no_payload)
return image
def info(f):
res = {}
m, handler = __rhandler(f)
res['magic'] = m
res['count'] = handler.count(f)
return res
def loads(s, pretty=False):
"""
Same as load(), but takes a string.
"""
f = io.BytesIO(s)
return load(f, pretty)
def dump(img, f):
"""
Convert criu image from dict(json) format to binary.
Takes an image in dict(json) format and file-like
object to write to.
"""
m = img['magic']
magic_val = magic.by_name[img['magic']]
# Images v1.1 NOTE: use "second" magic to identify what "first"
# should be written.
if m != 'INVENTORY':
if m in ('STATS', 'IRMAP_CACHE'):
f.write(struct.pack('i', magic.by_name['IMG_SERVICE']))
else:
f.write(struct.pack('i', magic.by_name['IMG_COMMON']))
f.write(struct.pack('i', magic_val))
try:
handler = handlers[m]
except:
raise Exception("No handler found for image with such magic")
handler.dump(img['entries'], f)
def dumps(img):
"""
Same as dump(), but takes only an image and returns
a string.
"""
f = io.BytesIO(b'')
dump(img, f)
return f.getvalue()
|