jw-devtest/src/python/devtest/os/be/qemu/Machine.py

# -*- coding: utf-8 -*-

import os
import asyncio
import subprocess
import importlib
import re
import io
import shutil
import functools
from jwutils.log import *
from jwutils.misc import get_derived_classes
from ...misc import *
from ...test import *
from ...Machine import Machine as MachineBase
from ...Connection import Connection
from ...conn.Fifos import Fifos as ConnFifos
from . import Invocation

_ifupdown_script="""#!/bin/bash

#echo exit for test purposes; exit 1
goodbye()
{
	:
	rm -rf $tmp_files
}

usage()
{
cat << EOT >&2

 $myname -h
 $myname {net|<image-file>} [-b hostname]

EOT
	[ "$1" ] && exit $1
	exit 0
}

log()
{
	local tag=`whoami`@$myname
	echo "$log_delim [$tag] $*"
	/usr/bin/logger -t "$tag" "$*"
}

err()
{
	log "$@" >&2
}

fatal()
{
	err "Fatal: $@ ... giving up"
	exit 1
}

run()
{
	log running $@
	"$@"
	return $?
}

bridge()
{
	case "$1" in
	start)
		die() {
			log $@
			bridge stop
			exit 1
		}

		failed_to() {
			die "failed to $@"
		}

		log setting up network on $bridge_name

		local -r net_prefix_len=$(echo $bridge_net | sed 's:.*/::')
		[ $net_prefix_len -lt 24 ] && fatal currently only class-C networks are supported
		local -r net_addr=$(echo $bridge_net | sed 's:/.*::')
		local -r last_octet_net_addr=$(echo $net_addr | sed 's/.*\.//')
		if [ "$bridge_ip" ]; then
			local_bridge_ip=$bridge_ip
		else
			local -r last_octet_bridge_ip=$(( $last_octet_net_addr + 1 ))
			local_bridge_ip=$(echo $net_addr | sed "s/\.[0-9]\+$/.$last_octet_bridge_ip/")
		fi
		local -r last_octet_broadcast=$(( $last_octet_net_addr + 2 ** (32 - $net_prefix_len) - 1 ))
		local -r broadcast=$(echo $net_addr | sed "s/\.[0-9]\+$/.$last_octet_broadcast/")

		run $brctl_exe addbr $bridge_name                     || failed_to "add bridge $bridge_name"
		#run $brctl_exe stp $bridge_name off                   || failed_to "disable spanning tree protocol on bridge $bridge_name"
		run $brctl_exe setfd $bridge_name 0                   || failed_to "set forward delay of bridge $bridge_name"
		run $ip_exe address add $local_bridge_ip broadcast $broadcast dev $bridge_name \
		                                                      || failed_to "add IP address to bridge $bridge_name"
		run $ip_exe link set $bridge_name up arp on           || failed_to "switch on arp on bridge $bridge_name"
		run $ip_exe route add to $bridge_net dev $bridge_name || failed_to "add route over bridge $bridge_name"
		;;
	stop)
		log shutting down network on $bridge_name
		run $ip_exe link set $bridge_name down
		run $brctl_exe delbr $bridge_name
		;;
	restart)
		bridge stop
		sleep 1
		bridge start
		;;
	check)
		LANG=POSIX $brctl_exe show $bridge_name 2>&1 | grep -q "$bridge_name.*No such device" && return 1
		$brctl_exe show $bridge_name >/dev/null 2>&1
		return $?
		;;
	*)
		echo "Usage: bridge {start|stop|restart|check}"
		exit 1
		;;
	esac
}

ref_bridge()
{
	bridge check || bridge start
}

unref_bridge()
{
	bridge check || return 0
	$brctl_exe show $bridge_name | awk "/^$bridge_name/ {print \$4}" | grep -q . || {
		log bridge $bridge_name is unused, cleaning up
		bridge stop
	}
	return 0
}

# -- here we go

myname=`basename $0`
log_delim="---------------"
log running $0 $@
[ -x $0 ] || chmod u+x $0 || fatal "$0 is not executable"
exe=`readlink -f $0`
dirname=`dirname $exe`
bridge_name=brdev0
#bridge_name=in1
base=dc=priv,dc=lcl
brctl_exe=`PATH=/usr/sbin:/sbin /usr/bin/which brctl`
ip_exe=`PATH=/usr/sbin:/sbin /usr/bin/which ip`
tmp_files=""
tmpdir=`dirname $0`
config=$tmpdir/config.sh

# -- qemu default options
bridge_net=192.168.100.0/24
mac_addr=,macaddr=00:0B:DC:9B:D6:DA

[ -r "$config" ] && . $config

trap goodbye INT QUIT EXIT KILL TERM PIPE

case $myname in
*ifup*)
	ref_bridge
	virt_if=$1
	run $ip_exe link set $virt_if up
	run $brctl_exe addif $bridge_name $virt_if
	;;

*ifdown*)
	virt_if=$1
	run $brctl_exe delif $bridge_name $virt_if
	run $ip_exe link set $virt_if down
	unref_bridge
	;;
*)
	echo called as unknown executable name $0
	;;
esac
"""

class SubprocessProtocol(asyncio.SubprocessProtocol):

    def __init__(self, machine, name):
        self.machine = machine
        self.name = name
        super().__init__()

    def pipe_data_received(self, fd, data):
        stream = "stdout" if fd == 1 else ("stderr" if fd == 2 else str(fd))
        tag = stream + '@' + self.name
        data = data.decode().rstrip('\n')
        prio = WARNING if fd == 2 else NOTICE
        for line in data.split('\n'):
            slog(prio, "[%s] %s" % (tag, line.rstrip('\r\n')))

    def process_exited(self):
        slog(NOTICE, "[%s] process exited" % (self.name))
        super().process_exited()
        self.machine.qemu_exited()

class Machine(MachineBase): # export

    def __init__(self, env):
        super().__init__(env)
        self.monitor = None
        self.console = None
        self.__running = False
        self.__shutdown_requested = False
        self.__clear_for_tests = True
        self.__net_helper = None
        self.__invocation = None
        self.__transport = None
        self.__protocol = None
        self.__proc = None
        self.__rc = None
        self.__task = None

    async def __await_monitor_prompt(self, act_timeout=2.):
        if await expect(self.monitor, regex='\(qemu\) ', subject="Qemu monitor prompt",
                act_timeout=act_timeout) is None:
            raise Exception("timed out waiting for Qemu monitor prompt")

    def __mkfifo(self, tp, out):
        name = self.fifo_base(tp) + "." + ("out" if out else "in")
        try:
            st = os.stat(name)
        except:
            slog(DEBUG, "running mkfifo(%s)" % name)
            os.mkfifo(name)
        return name

    def __mkfifos(self, tp):
        return (self.__mkfifo(tp, True), self.__mkfifo(tp, False))

    async def __exec_qemu(self, env):

        def format_cmdline(arr):
            r = ''
            for tok in arr:
                if re.search(' ', tok):
                    r += ' "%s"' % tok
                    continue
                r += ' ' + tok
            return r[1:]

        try:
            slog(INFO, "opening read-FIFOs to machine")
            for name in [ "monitor", "console" ]:
                c = ConnFifos(self.env, "Serial", paths=self.__mkfifos(name))
                await c.fifo(ConnFifos.Dir.In).open(timeout=1.)
                setattr(self, name, c)
            slog(NOTICE, "==== invoking qemu: ", format_cmdline(self.invocation.cmdline))
            self.__transport, self.__protocol = await env.eloop.subprocess_exec(
                lambda: SubprocessProtocol(self, "qemu"),
                *self.invocation.cmdline,
            )
            self.__proc = self.__transport.get_extra_info('subprocess') # Popen instance

            for name in [ "monitor", "console" ]:
                c = self.__dict__[name]
                await c.fifo(ConnFifos.Dir.Out).open(retry=4, retry_log_level=INFO)
            await self.__await_monitor_prompt()
        except:
            slog(ERR, "failed to run Qemu process")
            raise

    def __reap_qemu(self):
        if self.__rc is None and self.__transport:
            self.__transport = None
            self.__rc = self.__proc.wait()

    async def __cleanup_qemu(self):
        pid = self.__reap_qemu()
        if self.__rc == 0 and self.__shutdown_requested:
            slog(NOTICE, "the Qemu process (pid {}) has exited cleanly".format(self.__proc.pid))
            self.monitor = self.console = self.__protocol = self.__task = None
            return 0
        self.__clear_for_tests = False
        slog(ERR, "the Qemu process (pid {}) has exited {}with status code {}, aborting test".format(
            pid, "" if self.__shutdown_requested else "prematurely ", self.__rc))
        exit(1)

    # ---- utilities

    # to be called from SIGCHLD handler
    def qemu_exited(self):
        slog(INFO, "Qemu process exited")
        self.__clear_for_tests = False
        self.__reap_qemu()
        #self.__cleanup_qemu()

    @property
    def invocation(self):
        if not self.__invocation:
            mod_name = "Invocation_" + self.env.args.platform
            mod = importlib.import_module("devtest.os.be.qemu." + mod_name)
            classes = get_derived_classes(mod, Invocation.Invocation)
            if not len(classes):
                raise Exception("unsupported platform >" + self.env.args.platform + "<")
            self.__invocation = classes[0](self)
            slog(DEBUG, "found invocation definition class >%s<" % self.__invocation.typename())
        return self.__invocation

    def net_helper_path(self, up):
        if re.search('ifup\|ifdown', self.exe_basename):
            return None
        if not self.__net_helper:
            slog_m(INFO, "config = >%s<" % self.invocation.helper_config)
            self.__net_helper = []
            script_name = 'ifupdown.sh'
            script_path = self.tmpdir + "/" + script_name
            if self.env.args.qe_network_script:
                shutil.copyfile(self.env.args.qe_network_script, script_path)
            else:
                with open(script_path, "w") as fd:
                    fd.write(_ifupdown_script)
            os.chmod(script_path, 0o755)
            with open(self.tmpdir + "/config.sh", "w") as fd:
                fd.write(self.invocation.helper_config)
            for name in [ "ifdown", "ifup" ]:
                path = "{}/{}.sh".format(self.tmpdir, name)
                os.symlink(script_name, path)
                self.__net_helper.append(path)
        return self.__net_helper[up]

    def fifo_base(self, tp):
        return self.tmpdir + "/qemu-" + self.env.args.platform + "-" + tp

    # ---- reimplementation of class API methods

    async def init(self):
        self.__task = await self.env.eloop.create_task(self.__exec_qemu(self.env))

    async def cleanup(self, env):
        await self.__cleanup_qemu()

    def clear_for_tests(self):
        return self.__clear_for_tests

    def register_connection(self, env, info):
        slog(NOTICE, 'registering Qemu connection "{}" to platform "{}"'.format(
            info, self.env.args.platform))
        try:
            if info.proto == Connection.Proto.Console:
                r = self.console # TODO: info.spec ignored
            else:
                return super().register_connection(env, info)
            self._register_connection(r)
        except Exception as e:
            slog(ERR, "failed to open connection {}: {}".format(info, e))
            raise
        return r

    async def unregister_connection(self, conn):
        return await super().unregister_connection(conn)

    async def request_power_on(self, env):
        if self.__running:
            raise Exception("Tried to power on a running Qemu machine")
        slog(NOTICE, "switching on DUT")
        await self.monitor.write(b'cont\n')
        await self.__await_monitor_prompt()
        slog(NOTICE, "switched on DUT")
        self.__running = True

    async def wait_up(self, env):
        return True

    async def request_shutdown(self, env):
        if not self.__shutdown_requested:
            slog(NOTICE, "requesting shutdown")
            self.__shutdown_requested = True
            if self.monitor and self.__rc is None:
                await self.monitor.write(b'quit\n')

    async def wait_poweroff(self, env):
        slog(NOTICE, "waiting on powerdown")
        await self.__cleanup_qemu()
        return True