criu/test/zdtm/static/netns_lock.hook

#!/usr/bin/env python

import subprocess
import socket
import sys
import os

from ctypes import CDLL

libc = CDLL('libc.so.6')

SO_MARK = 36 # Define socket option for python2 compatibility
SOCCR_MARK = 0xC114
CLONE_NEWNET = 0x40000000
PORT = 8880
NETNS = "criu-net-lock-test"
TIMEOUT = 0.1

def nsenter():
    with open("/var/run/netns/{}".format(NETNS)) as f:
        libc.setns(f.fileno(), CLONE_NEWNET)

if sys.argv[1] == "--post-start":
    # Add test netns
    subprocess.Popen(["ip", "netns", "add", NETNS]).wait()
    nsenter() # Enter test netns
    subprocess.Popen(["ip", "link", "set", "up", "dev", "lo"]).wait()

    # TCP server daemon
    pid = os.fork()
    if(pid == 0):
        os.setsid() # Detach from parent
        # Run server
        srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        srv.bind(("127.0.0.1", PORT))
        srv.listen(1)

        # We should accept connections normally from
        # pre-dump client, post-restore client and
        # the pre-restore client using SOCCR_MARK
        # The pre-restore client without SOCCR_MARK
        # should fail with a timeout

        # Accept pre-dump client
        cln, addr = srv.accept()
        cln.sendall(str.encode("--pre-dump"))
        cln.close()

        # Accept pre-restore client with SOCCR_MARK
        srv.setsockopt(socket.SOL_SOCKET, SO_MARK, SOCCR_MARK)
        cln, addr = srv.accept()
        cln.sendall(str.encode("--pre-restore"))
        cln.close()
        srv.setsockopt(socket.SOL_SOCKET, SO_MARK, 0)

        # Accept post-restore client
        cln, addr = srv.accept()
        cln.sendall(str.encode("--post-restore"))
        cln.close()

        srv.close()

if sys.argv[1] == "--pre-dump":
    # Network is not locked yet
    # Client should be able to connect normally

    nsenter() # Enter test netns

    socket.setdefaulttimeout(TIMEOUT)
    cln = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        cln.connect(("127.0.0.1", PORT))
        resp = cln.recv(100).decode()
        if resp != sys.argv[1]:
            print("Expected '{}', got '{}'".format(sys.argv[1], resp))
            sys.exit(1)
    except socket.timeout:
        print("Timeout when trying to connect to server")
        sys.exit(1)
    else:
        sys.exit(0)
    finally:
        cln.close()

if sys.argv[1] == "--pre-restore":
    # Network should be locked now
    # Only packets with SOCCR_MARK should be allowed
    # Client should timeout when connecting without SOCCR_MARK

    nsenter()

    socket.setdefaulttimeout(TIMEOUT)
    # SOCCR_MARK
    try:
        cln = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        cln.setsockopt(socket.SOL_SOCKET, SO_MARK, SOCCR_MARK)
        cln.connect(("localhost", PORT))
        resp = cln.recv(100).decode()
        if resp != sys.argv[1]:
            print("Expected '{}', got '{}'".format(sys.argv[1], resp))
            sys.exit(1)
    except socket.timeout:
        print("Timeout when trying to connect to server")
        sys.exit(1)
    finally:
        cln.close()

    # No SOCCR_MARK
    try:
        cln = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        cln.connect(("localhost", PORT))
    except socket.timeout:
        sys.exit(0)
    else:
        print("Network should be locked, but client connected normally")
        sys.exit(1)
    finally:
        cln.close()

if sys.argv[1] == "--post-restore":
    # Network should be unlocked now
    # Client should be able to connect normally

    nsenter()

    socket.setdefaulttimeout(TIMEOUT)
    cln = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        cln.connect(("localhost", PORT))
        resp = cln.recv(100).decode()
        if resp != sys.argv[1]:
            print("Expected '{}', got '{}'".format(sys.argv[1], resp))
            sys.exit(1)
    except socket.timeout:
        print("Timeout when trying to connect to server")
        sys.exit(1)
    else:
        sys.exit(0)
    finally:
        cln.close()

if sys.argv[1] == "--clean":
    # Delete test netns
    subprocess.Popen(["ip", "netns", "delete", NETNS]).wait()