xen/domUloader.py

#!/usr/bin/env python
# domUloader.py
"""Loader for kernel and (optional) ramdisk from domU filesystem
 
 Uses bootentry = [dev:]kernel[,initrd] to get a kernel [and initrd]
  from a domU filesystem to boot it for Xen.
  dev is the disk as seen by domU, filenames are relative to
  that filesystem. The script uses the disk settings from the
  config file to find the domU filesystems.
  The bootentry is passed to the script using --entry=
 Optionally, dev: can be omitted; the script then looks at the
  root filesystem, parses /etc/fstab to resolve the path to the
  kernel [and the initrd]. Obviously, the paths relative to the
  domU root filesystem needs to be specified for the kernel
  and initrd filenames.

 The root FS is passed using --root=, the filesystem setup in
  --disks=. The disks list is a python list 
  [[uname, dev, mode, backend], [uname, dev, mode, backend], ...]
  passed as a string. The script writes an sxpr specifying the
  locations of the copied kernel and initrd into the file
  specified by --output (default is stdout).

 Limitations:
 - It is assumed both kernel and initrd are on the same filesystem.
 - domUs might use LVM; the script currently does not have support
   for setting up LVM mappings for domUs; it's not trivial and we 
   might risk namespace conflicts. If you want to use LVM inside domUs,
   set up a small non-LVM boot partition and specify it in bootentry.
 
 The script uses kpartx (multipath-tools) to create mappings for 
 devices that are exported as whole disk devices that are partitioned.

 (c) 01/2006 Novell Inc
 License: GNU GPL
 Author: Kurt Garloff <garloff@suse.de>
"""

import os, sys, getopt
from stat import *
from xen.xend import sxp
import tempfile
import time

# Global options
quiet = False
dryrun = False
verbose = False
tmpdir = '/var/lib/xen/tmp'

# List of partitions
# It's created by setting up the all the devices from the xen disk
# config; every entry creates on Wholedisk object, which does necessary
# preparatory steps such as losetup and kpartx -a; then a Partition
# object is setup for every partition (which may be one or several per
# Wholedisk); it references the Wholedisk if needed; python reference
# counting will take care of the cleanup.
partitions = []

# Helper functions

def traildigits(strg):
    "Return the trailing digits, used to split the partition number off"
    idx = len(strg)-1
    while strg[idx].isdigit():
        if len == 0:
            return strg
        idx -= 1
    return strg[idx+1:]    

def isWholedisk(domUname):
    "Determines whether dev is a wholedisk dev"
    return not domUname[-1:].isdigit()


def findPart(dev):
    "Find device dev in list of partitions"
    if len(dev) > 5 and dev[:5] == "/dev/":
        dev = dev[5:]
    for part in partitions:
        if dev == part.domname:
            return part
    return None

class Wholedisk:
    "Class representing a whole disk that has partitions"
    def    __init__(self, domname, physdev, loopfile = None):
        "c'tor: set up"
        self.domname  = domname
        self.physdev  = physdev
        self.loopfile = loopfile
        self.mapped   = 0
        self.pcount   = self.scanpartitions()

    def loopsetup(self):
        "Setup the loop mapping"
        if self.loopfile and not self.physdev:
            # Loops through all loopback devices, attempting to
            # find a free one to set up.  Don't scan for free and
            # then try to set it up as a separate step - too racy!
            i = 0
            while True:
                ldev = '/dev/loop%i' % (i)
                if not os.path.exists(ldev):
                    break
                i += 1
                fd = os.popen("losetup %s %s 2> /dev/null" % (ldev, self.loopfile))
                if not fd.close():
                    if verbose:
                        print "domUloader: losetup %s %s" % (ldev, self.loopfile)
                    self.physdev = ldev
                    break
            if not self.physdev:
                raise RuntimeError("domUloader: No free loop device found")

    def loopclean(self):
        "Delete the loop mapping"
        if self.loopfile and self.physdev:
            if verbose:
                print "domUloader: losetup -d %s" % self.physdev
            # Even seemingly innocent queries like "losetup /dev/loop0"
            # can temporarily block the loopback and cause transient
            # failures deleting the loopback, hence the retry logic.
            retries = 10
            while retries:
                fd = os.popen("losetup -d %s" % self.physdev)
                if not fd.close():
                    self.physdev = None
                    break
                else:
                    time.sleep(0.1)
                    retries -= 1
    
    def scanpartitions(self):
        """Scan device for partitions (kpartx -l) and set up data structures,
           Returns number of partitions found."""
        self.loopsetup()
        # TODO: We could use fdisk -l instead and look at the type of
        #  partitions; this way we could also detect LVM and support it.
        fd = os.popen("kpartx -l %s" % self.physdev)
        pcount = 0
        for line in fd.readlines():
            line = line.strip()
            (pname, params) = line.split(':')
            pno = int(traildigits(pname.strip()))
            #if pname.rfind('/') != -1:
            #    pname = pname[pname.rfind('/')+1:]
            #pname = self.physdev[:self.physdev.rfind('/')] + '/' + pname
            pname = "/dev/mapper/" + pname
            partitions.append(Partition(self, self.domname + '%i' % pno, pname))
            pcount += 1
        fd.close()
        if not pcount:
            if self.loopfile:
                ref = self
            else:
                ref = None
            partitions.append(Partition(ref, self.domname, self.physdev))
        return pcount

    def activatepartitions(self):
        "Set up loop mapping and device-mapper mappings"
        if not self.mapped:
            self.loopsetup()
            if self.pcount:
                if verbose:
                    print "domUloader: kpartx -a %s" % self.physdev
                fd = os.popen("kpartx -a %s" % self.physdev)
                fd.close()
        self.mapped += 1
    
    def deactivatepartitions(self):
        "Remove device-mapper mappings and loop mapping"
        if not self.mapped:
            return
        self.mapped -= 1
        if not self.mapped:
            if self.pcount:
                if verbose:
                    print "domUloader: kpartx -d %s" % self.physdev
                fd = os.popen("kpartx -d %s" % self.physdev)
                fd.close()
            self.loopclean()

    def __del__(self):
        "d'tor: clean up"
        self.deactivatepartitions()
        self.loopclean()

    def __repr__(self):
        "string representation for debugging"
        strg = "[" + self.domname + ","
        if self.physdev:
            strg += self.physdev
        strg += ","
        if self.loopfile:
            strg += self.loopfile
        strg += "," + str(self.pcount) + ",mapped %ix]" % self.mapped
        return strg

class Partition:
    """Class representing a domU filesystem (partition) that can be
       mounted in dom0"""
    def __init__(self, whole = None, domname = None, 
                 physdev = None):
        "c'tor: setup"
        self.wholedisk  = whole
        self.domname    = domname
        self.physdev    = physdev
        self.mountpoint = None

    def __del__(self):
        "d'tor: cleanup"
        if self.mountpoint:
            self.umount()
        # Not needed: Refcounting will take care of it.
        #if self.wholedisk:
        #    self.wholedisk.deactivatepartitions()

    def __repr__(self):
        "string representation for debugging"
        strg = "[" + self.domname + "," + self.physdev + ","
        if self.mountpoint:
            strg += "mounted on " + self.mountpoint + ","
        else:
            strg += "not mounted,"
        if self.wholedisk:
            return strg + self.wholedisk.__repr__() + "]"
        else:
            return strg + "]"

    def mount(self, fstype = None, options = "ro"):
        "mount filesystem, sets self.mountpoint"
        if self.mountpoint:
            return
        if self.wholedisk:
            self.wholedisk.activatepartitions()
        mtpt = tempfile.mkdtemp(prefix = "%s." % self.domname, dir = tmpdir)
        mopts = ""
        if fstype:
            mopts += " -t %s" % fstype
        mopts += " -o %s" % options
        if verbose:
            print "domUloader: mount %s %s %s" % (mopts, self.physdev, mtpt)
        fd = os.popen("mount %s %s %s" % (mopts, self.physdev, mtpt))
        err = fd.close()
        if err:
            raise RuntimeError("domUloader: Error %i from mount %s %s on %s" % \
                (err, mopts, self.physdev, mtpt))
        self.mountpoint = mtpt
        
    def umount(self):
        "umount filesystem at self.mountpoint"
        if not self.mountpoint:
            return
        if verbose:
            print "domUloader: umount %s" % self.mountpoint
        fd = os.popen("umount %s" % self.mountpoint)
        err = fd.close()
        os.rmdir(self.mountpoint)
        if err:
            raise RuntimeError("domUloader: Error %i from umount %s" % \
                (err, self.mountpoint))
        self.mountpoint = None
        if self.wholedisk:
            self.wholedisk.deactivatepartitions()


def setupOneDisk(cfg):
    """Sets up one exported disk (incl. partitions if existing)
    @param cfg: 4-tuple (uname, dev, mode, backend)""" 
    from xen.util.blkif import blkdev_uname_to_file
    values = cfg[0].split(':')
    if len(values) == 2:
        (type, dev) = values
    else:
        (type, subtype, dev) = values
    (loopfile, physdev) = (None, None)
    if type == "tap":
        if subtype == "aio":
            # FIXME:  if device, "/dev/" may need to be prepended
            mode = os.stat(dev)[ST_MODE]
            if S_ISBLK(mode):
                physdev = dev
            else:
                loopfile = dev
        else:
            raise RuntimeError("domUloader: domUloader supports 'tap' only with 'aio'.")
    if type == "file":
        loopfile = dev
    elif type == "phy":
        physdev = blkdev_uname_to_file(cfg[0])
    wdisk = Wholedisk(cfg[1], physdev, loopfile)

def setupDisks(vbds):
    """Create a list of all disks from the disk config:
    @param vbds: The disk config as list of 4-tuples 
                 (uname, dev, mode, backend)"""
    disks = eval(vbds)
    for disk in disks:
        setupOneDisk(disk)
    if verbose:
        print "Partitions: " + str(partitions)

class Fstab:
    "Class representing an fstab"

    class FstabEntry:
        "Class representing one fstab line"
        def __init__(self, line):
            "c'tor: parses one line"
            spline = line.split()
            self.dev, self.mtpt, self.fstype, self.opts = \
                spline[0], spline[1], spline[2], spline[3]
            if len(self.mtpt) > 1:
                self.mtpt = self.mtpt.rstrip('/')
    
    def __init__(self, filename):
        "c'tor: parses fstab"
        self.entries = []
        fd = open(filename)
        for line in fd.readlines():
            line = line.strip()
            if len(line) == 0 or line[0] == '#':
                continue
            self.entries.append(Fstab.FstabEntry(line))

    def find(self, fname):
        "Looks for matching filesystem in fstab"
        matchlen = 0
        match = None
        fnmlst = fname.split('/')
        for fs in self.entries:
            entlst = fs.mtpt.split('/')
            # '/' needs special treatment :-(
            if entlst == ['','']:
                entlst = ['']
            entln = len(entlst)
            if len(fnmlst) >= entln and fnmlst[:entln] == entlst \
                and entln > matchlen:
                match = fs
                matchlen = entln
        if not match:
            return (None, None)
        return (match.dev, match.mtpt)

def fsFromFstab(kernel, initrd, root):
    """Investigate rootFS fstab, check for filesystem that contains the kernel
    and return it; also returns adapted kernel and initrd path.
    """
    part = findPart(root)
    if not part:
        raise RuntimeError("domUloader: Root fs %s not exported?" % root)
    part.mount()
    if not os.access(part.mountpoint + '/etc/fstab', os.R_OK):
        part.umount()
        raise RuntimeError("domUloader: /etc/fstab not found on %s" % root)
    fstab = Fstab(part.mountpoint + '/etc/fstab')
    (dev, fs) = fstab.find(kernel)
    if not fs:
        raise RuntimeError("domUloader: no matching filesystem for image %s found in fstab" % kernel)
        #return (None, kernel, initrd)
    if fs == '/':
        ln = 0
        # this avoids the stupid /dev/root problem
        dev = root
    else:
        ln = len(fs)
    kernel = kernel[ln:]
    if initrd:
        initrd = initrd[ln:]
    if verbose:
        print "fsFromFstab: %s %s -- %s,%s" % (dev, fs, kernel, initrd)
    return (kernel, initrd, dev) 

def parseEntry(entry):
    "disects bootentry and returns kernel, initrd, filesys"
    fs = None
    initrd = None
    fsspl = entry.split(':')
    if len(fsspl) > 1:
        fs = fsspl[0]
        entry = fsspl[1]
    enspl = entry.split(',')
    # Prepend '/' if missing
    kernel = enspl[0]
    if kernel[0] != '/':
        kernel = '/' + kernel
    if len(enspl) > 1:
        initrd = enspl[1]
        if initrd[0] != '/':
            initrd = '/' + initrd
    return kernel, initrd, fs

def copyFile(src, dst):
    "Wrapper for shutil.filecopy"
    import shutil
    if verbose:
        print "domUloader: cp %s %s" % (src, dst)
    stat = os.stat(src)
    if stat.st_size > 16*1024*1024:
        raise RuntimeError("domUloader: Too large file %s (%s larger than 16MB)" \
            % (src, stat.st_size)) 
    try:
        shutil.copyfile(src, dst)
    except:
        os.unlink(dst)
        raise()

def copyKernelAndInitrd(fs, kernel, initrd):
    """Finds fs in list of partitions, mounts the partition, copies
       kernel [and initrd] off to dom0 files, umounts the parition again,
       and returns sxpr pointing to these copies."""
    if dryrun:
        return "linux (kernel vmlinuz.dummy) (ramdisk initrd.dummy)"
    import shutil
    part = findPart(fs)
    if not part:
        raise RuntimeError("domUloader: Filesystem %s not exported\n" % fs)
    part.mount()
    try:
        (fd, knm) = tempfile.mkstemp(prefix = "vmlinuz.", dir = tmpdir)
        os.close(fd)
        copyFile(part.mountpoint + kernel, knm)
    except: 
        os.unlink(knm)
        raise
    if not quiet:
        print "Copy kernel %s from %s to %s for booting" % \
            (kernel, fs, knm)
    sxpr = "linux (kernel %s)" % knm
    if (initrd):
        try:
            (fd, inm) = tempfile.mkstemp(prefix = "initrd.", dir = tmpdir)
            os.close(fd)
            copyFile(part.mountpoint + initrd, inm)
        except:
            os.unlink(knm)
            os.unlink(inm)
            raise
        sxpr += "(ramdisk %s)" % inm
    part.umount()
    return sxpr

def main(argv):
    "Main routine: Parses options etc."
    global quiet, dryrun, verbose, tmpdir
    def usage():
        "Help output (usage info)"
        global verbose, quiet, dryrun
        print >> sys.stderr, "domUloader usage: domUloader --disks=disklist [--root=rootFS]\n" \
            + " --entry=kernel[,initrd] [--output=fd] [--quiet] [--dryrun] [--verbose] [--help]\n"
        print >> sys.stderr, __doc__

    #print "domUloader " + str(argv)
    try:
        (optlist, args) = getopt.gnu_getopt(argv, 'qvh', \
            ('disks=', 'root=', 'entry=', 'output=', 
             'tmpdir=', 'help', 'quiet', 'dryrun', 'verbose'))
    except:
        usage()
        sys.exit(1)

    entry = None
    output = None
    root = None
    disks = None

    for (opt, oarg) in optlist:
        if opt in ('-h', '--help'):
            usage()
            sys.exit(0)
        elif opt in ('-q', '--quiet'):
            quiet = True
        elif opt in ('-n', '--dryrun'):
            dryrun = True
        elif opt in ('-v', '--verbose'):
            verbose = True
        elif opt == '--root':
            root = oarg
        elif opt == '--output':
            output = oarg
        elif opt == '--disks':
            disks = oarg
        elif opt == '--entry':
            entry = oarg
        elif opt == '--tmpdir':
            tmpdir = oarg

    if not entry or not disks:
        usage()
        sys.exit(1)
    
    if output is None or output == "-":
        fd = sys.stdout.fileno()
    else:
        fd = os.open(output, os.O_WRONLY)
    
    if not os.access(tmpdir, os.X_OK):
        os.mkdir(tmpdir)
        os.chmod(tmpdir, 0750)

    # We assume kernel and initrd are on the same FS,
    # so only one fs
    kernel, initrd, fs = parseEntry(entry)
    setupDisks(disks)
    if not fs:
        if not root:
            usage()
            raise RuntimeError("domUloader: No root= to parse fstab and no disk in bootentry")
            sys.exit(1)
        kernel, initrd, fs = fsFromFstab(kernel, initrd, root)
    
    sxpr = copyKernelAndInitrd(fs, kernel, initrd)
    sys.stdout.flush()
    os.write(fd, sxpr)

# Call main if called (and not imported)
if __name__ == "__main__":
    main(sys.argv)