diff --git a/pycriu/__init__.py b/pycriu/__init__.py index 52059738f..7de62838c 100644 --- a/pycriu/__init__.py +++ b/pycriu/__init__.py @@ -1,2 +1,3 @@ import rpc import images +from criu import * diff --git a/pycriu/criu.py b/pycriu/criu.py new file mode 100644 index 000000000..84dcefed2 --- /dev/null +++ b/pycriu/criu.py @@ -0,0 +1,282 @@ +# Same as libcriu for C. + +import socket +import errno +import subprocess +import fcntl +import os +import signal +import sys +import struct + +import rpc + +class _criu_comm: + """ + Base class for communication classes. + """ + COMM_SK = 0 + COMM_FD = 1 + COMM_BIN = 2 + comm_type = None + comm = None + sk = None + + def connect(self, daemon): + """ + Connect to criu and return socket object. + daemon -- is for whether or not criu should daemonize if executing criu from binary(comm_bin). + """ + pass + + def disconnect(self): + """ + Disconnect from criu. + """ + pass + + +class _criu_comm_sk(_criu_comm): + """ + Communication class for unix socket. + """ + def __init__(self, sk_path): + self.comm_type = self.COMM_SK + self.comm = sk_path + + def connect(self, daemon): + self.sk = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET) + self.sk.connect(self.comm) + + return self.sk + + def disconnect(self): + self.sk.close() + + +class _criu_comm_fd(_criu_comm): + """ + Commnunication class for file descriptor. + """ + def __init__(self, fd): + self.comm_type = self.COMM_FD + self.comm = fd + + def connect(self, daemon): + self.sk = socket.fromfd(self.comm, socket.AF_UNIX, socket.SOCK_SEQPACKET) + + return self.sk + + def disconnect(self): + self.sk.close() + +class _criu_comm_bin(_criu_comm): + """ + Communication class for binary. + """ + def __init__(self, bin_path): + self.comm_type = self.COMM_BIN + self.comm = bin_path + self.swrk = None + self.daemon = None + + def connect(self, daemon): + # Kind of the same thing we do in libcriu + css = socket.socketpair(socket.AF_UNIX, socket.SOCK_SEQPACKET) + flags = fcntl.fcntl(css[1], fcntl.F_GETFD) + fcntl.fcntl(css[1], fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC) + + self.daemon = daemon + + p = os.fork() + + if p == 0: + def exec_criu(): + os.close(0) + os.close(1) + os.close(2) + + css[0].send(struct.pack('i', os.getpid())) + os.execv(self.comm, [self.comm, 'swrk', "%d" % css[0].fileno()]) + os._exit(1) + + if daemon: + # Python has no daemon(3) alternative, + # so we need to mimic it ourself. + p = os.fork() + + if p == 0: + os.setsid() + + exec_criu() + else: + os._exit(0) + else: + exec_criu() + + css[0].close() + self.swrk = struct.unpack('i', css[1].recv(4))[0] + self.sk = css[1] + + return self.sk + + def disconnect(self): + self.sk.close() + if not self.daemon: + os.waitpid(self.swrk, 0) + + +class CRIUException(Exception): + """ + Exception class for handling and storing criu errors. + """ + typ = None + _str = None + + def __str__(self): + return self._str + + +class CRIUExceptionInternal(CRIUException): + """ + Exception class for handling and storing internal errors. + """ + def __init__(self, typ, s): + self.typ = typ + self._str = "%s failed with internal error: %s" % (rpc.criu_req_type.Name(self.typ), s) + + +class CRIUExceptionExternal(CRIUException): + """ + Exception class for handling and storing criu RPC errors. + """ + + def __init__(self, req_typ, resp_typ, errno): + self.typ = req_typ + self.resp_typ = resp_typ + self.errno = errno + self._str = self._gen_error_str() + + def _gen_error_str(self): + s = "%s failed: " % (rpc.criu_req_type.Name(self.typ), ) + + if self.typ != self.resp_typ: + s += "Unxecpected response type %d: " % (self.resp_typ, ) + + s += "Error(%d): " % (self.errno, ) + + if self.errno == errno.EBADRQC: + s += "Bad options" + + if self.typ == rpc.DUMP: + if self.errno == errno.ESRCH: + s += "No process with such pid" + + if self.typ == rpc.RESTORE: + if self.errno == errno.EEXIST: + s += "Process with requested pid already exists" + + s += "Unknown" + + return s + + +class criu: + """ + Call criu through RPC. + """ + opts = None #CRIU options in pb format + + _comm = None #Communication method + + def __init__(self): + self.use_binary('criu') + self.opts = rpc.criu_opts() + + def use_sk(self, sk_name): + """ + Access criu using unix socket which that belongs to criu service daemon. + """ + self._comm = _criu_comm_sk(sk_name) + + def use_fd(self, fd): + """ + Access criu using provided fd. + """ + self._comm = _criu_comm_fd(fd) + + def use_binary(self, bin_name): + """ + Access criu by execing it using provided path to criu binary. + """ + self._comm = _criu_comm_bin(bin_name) + + def _send_req_and_recv_resp(self, req): + """ + As simple as send request and receive response. + """ + # In case of self-dump we need to spawn criu swrk detached + # from our current process, as criu has a hard time separating + # process resources from its own if criu is located in a same + # process tree it is trying to dump. + daemon = False + if req.type == rpc.DUMP and not req.opts.HasField('pid'): + daemon = True + + try: + s = self._comm.connect(daemon) + + s.send(req.SerializeToString()) + + buf = s.recv(len(s.recv(1, socket.MSG_TRUNC | socket.MSG_PEEK))) + + self._comm.disconnect() + + resp = rpc.criu_resp() + resp.ParseFromString(buf) + except Exception as e: + raise CRIUExceptionInternal(req.type, str(e)) + + return resp + + def check(self): + """ + Checks whether the kernel support is up-to-date. + """ + req = rpc.criu_req() + req.type = rpc.CHECK + + resp = self._send_req_and_recv_resp(req) + + if not resp.success: + raise CRIUExceptionExternal(req.type, resp.type, resp.cr_errno) + + def dump(self): + """ + Checkpoint a process/tree identified by opts.pid. + """ + req = rpc.criu_req() + req.type = rpc.DUMP + req.opts.MergeFrom(self.opts) + + resp = self._send_req_and_recv_resp(req) + + if not resp.success: + raise CRIUExceptionExternal(req.type, resp.type, resp.cr_errno) + + return resp.dump + + def restore(self): + """ + Restore a process/tree. + """ + req = rpc.criu_req() + req.type = rpc.RESTORE + req.opts.MergeFrom(self.opts) + + resp = self._send_req_and_recv_resp(req) + + if not resp.success: + raise CRIUExceptionExternal(req.type, resp.type, resp.errno) + + return resp.restore