criu-ns: Add tests for criu-ns script

These changes add test implementations for criu-ns script.

Fixes: checkpoint-restore#1909

Signed-off-by: Dhanuka Warusadura <csx@tuta.io>

Makefile

@@ -438,6 +438,7 @@ lint:
 	flake8 --config=scripts/flake8.cfg lib/py/images/pb2dict.py
 	flake8 --config=scripts/flake8.cfg lib/py/images/images.py
 	flake8 --config=scripts/flake8.cfg scripts/criu-ns
+	flake8 --config=scripts/flake8.cfg test/others/criu-ns/run.py
 	flake8 --config=scripts/flake8.cfg crit/setup.py
 	flake8 --config=scripts/flake8.cfg scripts/uninstall_module.py
 	flake8 --config=scripts/flake8.cfg coredump/

scripts/ci/run-ci-tests.sh

@@ -259,6 +259,7 @@ if [ -n "$TRAVIS" ] || [ -n "$CIRCLECI" ]; then
        # GitHub Actions (and Cirrus CI) does not provide a real TTY and CRIU will fail with:
        # Error (criu/tty.c:1014): tty: Don't have tty to inherit session from, aborting
        make -C test/others/shell-job/ run
+       make -C test/others/criu-ns/ run
 fi
 make -C test/others/skip-file-rwx-check/ run
 make -C test/others/rpc/ run

test/Makefile

@@ -12,7 +12,7 @@ all:
 	$(MAKE) zdtm-freezer
 .PHONY: all
 
-TESTS = unix-callback mem-snap rpc libcriu mounts/ext security pipes crit socketpairs overlayfs mnt-ext-dev shell-job skip-file-rwx-check
+TESTS = unix-callback mem-snap rpc libcriu mounts/ext security pipes crit socketpairs overlayfs mnt-ext-dev shell-job criu-ns skip-file-rwx-check
 
 other:
 	for t in $(TESTS); do				\

test/others/criu-ns/Makefile

@@ -0,0 +1,3 @@
+run:
+	@make -C ../.. zdtm_ct
+	../../zdtm_ct run.py

test/others/criu-ns/run.py

@@ -0,0 +1,182 @@
+#!/usr/bin/env python
+
+import os
+import pathlib
+import pty
+import shutil
+import subprocess
+import sys
+import time
+
+
+CRIU_BIN = "../../../criu/criu"
+CRIU_NS = "../../../scripts/criu-ns"
+IMG_DIR = "dumpdir"
+PIDFILE = "pidfile"
+
+
+def check_dumpdir(path=IMG_DIR):
+    if os.path.isdir(path):
+        shutil.rmtree(path)
+    os.mkdir(path, 0o755)
+
+
+def create_pty():
+    fd_m, fd_s = pty.openpty()
+    return (os.fdopen(fd_m, "wb"), os.fdopen(fd_s, "wb"))
+
+
+def create_isolated_dumpee():
+    pathlib.Path("running").touch()
+    fd_m, fd_s = create_pty()
+    pid = os.fork()
+    fd_m.close()
+    if pid == 0:
+        os.setsid()
+        os.dup2(fd_s.fileno(), 0)
+        os.dup2(fd_s.fileno(), 1)
+        os.dup2(fd_s.fileno(), 2)
+        while True:
+            if not os.access("running", os.F_OK):
+                sys.exit(0)
+            time.sleep(1)
+    return (pid, fd_s)
+
+
+def criu_ns_dump(pid, shell_job=False):
+    if shell_job:
+        cmd = [CRIU_NS, "dump", "-D", IMG_DIR, "-v4",
+               "--shell-job", "-t", str(pid), "--criu-binary",
+               CRIU_BIN]
+    else:
+        cmd = [CRIU_NS, "dump", "-D", IMG_DIR, "-v4",
+               "-t", str(pid), "--criu-binary", CRIU_BIN]
+    ret = subprocess.Popen(cmd).wait()
+    return ret
+
+
+def criu_ns_restore(shell_job=False, restore_detached=False):
+    if shell_job:
+        cmd = [CRIU_NS, "restore", "-D", IMG_DIR, "-v4",
+               "--shell-job", "--criu-binary", CRIU_BIN]
+    else:
+        cmd = [CRIU_NS, "restore", "-D", IMG_DIR, "-v4",
+               "--criu-binary", CRIU_BIN]
+    if restore_detached:
+        cmd = [CRIU_NS, "restore", "-D", IMG_DIR, "-v4",
+               "--restore-detached", "--pidfile", PIDFILE,
+               "--criu-binary", CRIU_BIN]
+    ret = subprocess.Popen(cmd).wait()
+    return ret
+
+
+def test_dump_and_restore_with_shell_job():
+    """Test criu-ns dump and restore with --shell-job option"""
+    check_dumpdir()
+    pathlib.Path("running").touch()
+    pid = os.fork()
+    if pid == 0:
+        while True:
+            if not os.access("running", os.F_OK):
+                sys.exit(0)
+            time.sleep(1)
+
+    ret = criu_ns_dump(pid, shell_job=True)
+    if ret != 0:
+        sys.exit(ret)
+
+    os.unlink("running")
+    fd_m, fd_s = create_pty()
+    pid = os.fork()
+    if pid == 0:
+        os.setsid()
+        fd_m.close()
+        # since criu-ns takes control of the tty stdin
+        os.dup2(fd_s.fileno(), 0)
+        ret = criu_ns_restore(shell_job=True)
+        if ret != 0:
+            sys.exit(ret)
+        os._exit(0)
+
+    fd_s.close()
+    os.waitpid(pid, 0)
+
+
+def test_dump_and_restore_without_shell_job(restore_detached=False):
+    """Test criu-ns dump and restore with an isolated process"""
+    check_dumpdir()
+    pid, fd_s = create_isolated_dumpee()
+    ret = criu_ns_dump(pid)
+    if ret != 0:
+        sys.exit(ret)
+
+    if not restore_detached:
+        os.unlink("running")
+
+    fd_m, fd_s = create_pty()
+    pid = os.fork()
+    if pid == 0:
+        os.setsid()
+        ret = criu_ns_restore(restore_detached=restore_detached)
+        if ret != 0:
+            sys.exit(ret)
+        os._exit(0)
+
+    fd_m.close()
+    fd_s.close()
+    os.waitpid(pid, 0)
+
+
+def test_dump_and_restore_in_pidns():
+    """Test criu-ns dump and restore in namespaces"""
+    def _dump():
+        pid, fd_s = create_isolated_dumpee()
+        ret = criu_ns_dump(pid)
+        if ret != 0:
+            sys.exit(ret)
+
+    def _restore():
+        ret = criu_ns_restore(restore_detached=True)
+        if ret != 0:
+            sys.exit(ret)
+
+    def _get_restored_pid():
+        restored_pid = 0
+        pidfile_path = os.path.join(IMG_DIR, PIDFILE)
+        if not os.path.exists(pidfile_path):
+            raise FileNotFoundError("pidfile not found")
+        with open(pidfile_path, "r") as pidfile:
+            restored_pid = pidfile.read().strip()
+        return int(restored_pid)
+
+    def _redump():
+        global IMG_DIR
+        try:
+            restored_pid = _get_restored_pid()
+        except FileNotFoundError:
+            sys.exit(1)
+        IMG_DIR = "dumpdir2"
+        check_dumpdir(IMG_DIR)
+        ret = criu_ns_dump(restored_pid)
+        if ret != 0:
+            sys.exit(ret)
+
+    def _re_restore():
+        os.unlink("running")
+        ret = criu_ns_restore()
+        if ret != 0:
+            sys.exit(ret)
+
+    check_dumpdir()
+    _dump()
+    _restore()
+    _redump()
+    _re_restore()
+
+
+if __name__ == "__main__":
+    os.setsid()
+    test_dump_and_restore_with_shell_job()
+    test_dump_and_restore_without_shell_job()
+    test_dump_and_restore_without_shell_job(restore_detached=True)
+    test_dump_and_restore_in_pidns()