#!/usr/bin/env python3
# Copyright 2017 Obsidian Research Corp.
# Licensed under BSD (MIT variant) or GPLv2. See COPYING.
"""check-build - Run static checks on a build"""
from __future__ import print_function
import argparse
import inspect
import os
import re
import shutil
import subprocess
import tempfile
import sys
import copy
import shlex
import pipes
from contextlib import contextmanager;
from pkg_resources import parse_version

def get_src_dir():
    """Get the source directory using git"""
    git_top = subprocess.check_output(["git","rev-parse","--git-dir"]).decode().strip();
    if git_top == ".git":
        return ".";
    return os.path.dirname(git_top);

def get_package_version(args):
    """Return PACKAGE_VERSION from CMake"""
    with open(os.path.join(args.SRC,"CMakeLists.txt")) as F:
        for ln in F:
            g = re.match(r'^set\(PACKAGE_VERSION "(.+)"\)',ln)
            if g is None:
                continue;
            return g.group(1);
    raise RuntimeError("Could not find version");

@contextmanager
def inDirectory(dir):
    cdir = os.getcwd();
    try:
        os.chdir(dir);
        yield True;
    finally:
        os.chdir(cdir);

@contextmanager
def private_tmp():
    """Simple version of Python 3's tempfile.TemporaryDirectory"""
    dfn = tempfile.mkdtemp();
    try:
        yield dfn;
    finally:
        shutil.rmtree(dfn);

# -------------------------------------------------------------------------

def get_symbol_vers(fn,exported=True):
    """Return the symbol version suffixes from the ELF file, eg IB_VERBS_1.0, etc"""
    syms = subprocess.check_output(["readelf","--wide","-s",fn]).decode();

    go = False;
    res = set();
    for I in syms.splitlines():
        if I.startswith("Symbol table '.dynsym'"):
            go = True;
            continue;

        if I.startswith(" ") and go:
            itms = I.split();
            if exported:
                if (len(itms) == 8 and itms[3] == "OBJECT" and
                    itms[4] == "GLOBAL" and itms[6] == "ABS"):
                    res.add(itms[7]);
            else:
                if (len(itms) >= 8 and itms[3] == "FUNC" and
                    itms[4] == "GLOBAL" and itms[6] == "UND"):
                    res.add(itms[7]);
        else:
            go = False;
    if not res:
        raise ValueError("Failed to read ELF symbol versions from %r"%(fn));
    return res;

def symver_parse_version(ver):
    return parse_version(ver.partition("_")[-1])

def check_lib_symver(args,fn):
    g = re.match(r"lib([^.]+)\.so\.(\d+)\.(\d+)\.(.*)",fn);
    if g.group(4) != args.PACKAGE_VERSION:
        raise ValueError("Shared Library filename %r does not have the package version %r (%r)"%(
            fn,args.PACKAGE_VERSION,g.groups()));

    # umad/etc used the wrong symbol version name when they moved to soname 3.0
    if g.group(1) == "ibumad":
        newest_symver = "%s_%s.%s"%(g.group(1).upper(),'1',g.group(3));
    elif g.group(1) == "ibmad":
        newest_symver = "%s_%s.%s"%(g.group(1).upper(),'1',g.group(3));
    elif g.group(1) == "ibnetdisc":
        newest_symver = "%s_%s.%s"%(g.group(1).upper(),'1',g.group(3));
    else:
        newest_symver = "%s_%s.%s"%(g.group(1).upper(),g.group(2),g.group(3));

    syms = get_symbol_vers(fn);
    if newest_symver not in syms:
        raise ValueError("Symbol version %r implied by filename %r not in ELF (%r)"%(
            newest_symver,fn,syms));

    # The private symbol tag should also be older than the package version
    private = set(I for I in syms if "PRIVATE" in I)
    if len(private) > 1:
        raise ValueError("Too many private symbol versions in ELF %r (%r)"%(fn,private));
    if private:
        private_rel = list(private)[0].split('_')[-1];
        if private_rel > args.PACKAGE_VERSION:
            raise ValueError("Private Symbol Version %r is newer than the package version %r"%(
                private,args.PACKAGE_VERSION));

    syms = list(syms - private);
    syms.sort(key=symver_parse_version)
    if newest_symver != syms[-1]:
        raise ValueError("Symbol version %r implied by filename %r not the newest in ELF (%r)"%(
            newest_symver,fn,syms));

def test_lib_names(args):
    """Check that the library filename matches the symbol versions"""
    libd = os.path.join(args.BUILD,"lib");

    # List of shlibs that follow the ABI guidelines
    libs = {};
    with inDirectory(libd):
        for fn in os.listdir("."):
            if os.path.islink(fn):
                lfn = os.readlink(fn);
                if not os.path.islink(lfn):
                    check_lib_symver(args,lfn);
# -------------------------------------------------------------------------

def check_verbs_abi(args,fn):
    g = re.match(r"lib([^-]+)-rdmav(\d+).so",fn);
    if g is None:
        raise ValueError("Provider library has unknown file name format %r"%(fn));

    private_ver = int(g.group(2));
    syms = get_symbol_vers(fn,exported=False);
    syms = {I.partition("@")[2] for I in syms};
    assert "IBVERBS_PRIVATE_%u"%(private_ver) in syms;
    assert len([I for I in syms if I.startswith("IBVERBS_PRIVATE")]) == 1;

def test_verbs_private(args):
    """Check that the IBVERBS_PRIVATE symbols match the library name, eg that the
    map file and the cmake stuff are in sync."""
    libd = os.path.join(args.BUILD,"lib");
    with inDirectory(libd):
        for fn in os.listdir("."):
            if not os.path.islink(fn) and "rdmav" in fn and fn.endswith(".so"):
                check_verbs_abi(args,fn);

# -------------------------------------------------------------------------

def check_abi(args,fn):
    g1 = re.match(r"lib([^.]+).so\.(.+)\.(.+)",fn);
    g2 = re.match(r"lib([^.]+).so\.(.+\..+)",fn);
    if g1 is None or g2 is None:
        raise ValueError("Library has unknown file name format %r"%(fn));

    ref_fn = os.path.join(args.SRC,"ABI",g1.group(1) + ".dump");
    cur_fn = os.path.join(args.SRC,"ABI","current-" + g1.group(1) + ".dump");
    subprocess.check_call(["abi-dumper",
                           "-lver",g2.group(1),
                           fn,
                           "-o",cur_fn]);

    if not os.path.exists(ref_fn):
        print("ABI file does not exist for %r"%(ref_fn), file=sys.stderr);
        return False;

    subprocess.check_call(["abi-compliance-checker",
                           "-l",g1.group(1),
                           "-old",ref_fn,
                           "-new",cur_fn]);

    return True;

def test_verbs_uapi(args):
    """Compare the ABI output from 'abi-dumper' between what is present in git and
    what was built in this tree.  This allows us to detect changes in ABI on
    the -stable branch."""

    # User must provide the ABI dir in the source tree
    if not os.path.isdir(os.path.join(args.SRC,"ABI")):
        print("ABI check skipped, no ABI/ directory.");
        return;

    libd = os.path.join(args.BUILD,"lib");
    success = True;
    with inDirectory(libd):
        for fn in os.listdir("."):
            if not os.path.islink(fn) and re.match(r"lib.+\.so\..+\..+",fn):
                success = success & check_abi(args,fn);

    assert success == True;

# -------------------------------------------------------------------------

def is_obsolete(fn):
    """True if the header is obsolete and should not be compiled anyhow."""
    with open(fn) as F:
        for ln in F.readlines():
            if re.search(r"#warning.*This header is obsolete",ln):
                return True;
    return False;

def is_fixup(fn):
    """True if this is a fixup header, fixup headers are exempted because they
    required includes are not the same for kernel headers (eg netinet/in.h)"""
    if os.path.islink(fn):
        return "buildlib/fixup-include/" in os.readlink(fn);
    return False;

def get_headers(incdir):
    includes = set();
    for root,dirs,files in os.walk(incdir):
        for I in files:
            if I.endswith(".h"):
                includes.add(os.path.join(root,I));
    return includes;

def compile_test_headers(tmpd,incdir,includes,with_cxx=False):
    cppflags = subprocess.check_output(["pkg-config","libnl-3.0","--cflags-only-I"]).decode().strip();
    cppflags = "-I %s %s"%(incdir,cppflags)
    with open(os.path.join(tmpd,"build.ninja"),"wt") as F:
        print("rule comp", file=F);
        print(" command = %s -Werror -c %s $in -o $out"%(args.CC,cppflags), file=F);
        print(" description=Header check for $in", file=F);
        print("rule comp_cxx", file=F);
        print(" command = %s -Werror -c %s $in -o $out"%(args.CXX,cppflags), file=F);
        print(" description=Header C++ check for $in", file=F);
        count = 0;
        for I in sorted(includes):
            if is_obsolete(I) or is_fixup(I):
                continue;
            print("build %s : comp %s"%("out%d.o"%(count),I), file=F);
            print("default %s"%("out%d.o"%(count)), file=F);
            print("build %s : comp_cxx %s"%("outxx%d.o"%(count),I), file=F);
            if with_cxx:
                print("default %s"%("outxx%d.o"%(count)), file=F);
            count = count + 1;
    subprocess.check_call(["ninja"],cwd=tmpd);

def test_published_headers(args):
    """Test that every header file can be included on its own, and has no obvious
    implicit dependencies. This is intended as a first pass check of the public
    installed API headers"""
    incdir = os.path.abspath(os.path.join(args.BUILD,"include"));
    includes = get_headers(incdir);

    # Make a little ninja file to compile each header
    with private_tmp() as tmpd:
        compile_test_headers(tmpd,incdir,includes);

# -------------------------------------------------------------------------

allowed_uapi_headers = {
    # This header is installed in all supported distributions
    "rdma/ib_user_sa.h",
    "rdma/ib_user_verbs.h",
    "linux/stddef.h",
}

non_cxx_headers = {
    "infiniband/arch.h",
    "infiniband/ib.h",
    "infiniband/ib_user_ioctl_verbs.h",
    "infiniband/ibnetdisc_osd.h",
    "infiniband/mad_osd.h",
    "infiniband/mlx5_api.h",
    "infiniband/mlx5_user_ioctl_verbs.h",
    "infiniband/opcode.h",
    "infiniband/sa-kern-abi.h",
    "infiniband/sa.h",
    "infiniband/verbs_api.h",
    "rdma/rdma_cma_abi.h",
}

def test_installed_headers(args):
    """This test also checks that the public headers can be compiled on their own,
    but goes further and confirms that the public headers do not depend on any
    internal headers, or kernel kAPI headers."""
    with private_tmp() as tmpd:
        env = copy.deepcopy(os.environ);
        env["DESTDIR"] = tmpd;
        subprocess.check_output(["ninja","install"],env=env,cwd=args.BUILD);

        includes = get_headers(tmpd);
        incdir = os.path.commonprefix(list(includes));
        rincludes = {I[len(incdir):] for I in includes};

        bincdir = os.path.abspath(os.path.join(args.BUILD,"include"));
        all_includes = set();
        for I in get_headers(bincdir):
            if not is_fixup(I) and not is_obsolete(I):
                all_includes.add(I[len(bincdir)+1:]);

        # Drop error includes for any include file that is internal, this way
        # when we compile the public headers any include of an internal header
        # will fail.
        for I in sorted(all_includes - rincludes):
            if I in allowed_uapi_headers:
                continue;

            I = os.path.join(incdir,I)
            dfn = os.path.dirname(I);
            if not os.path.isdir(dfn):
                os.makedirs(dfn);
            assert not os.path.exists(I);
            with open(I,"w") as F:
                print('#error "Private internal header"', file=F);

        # Roughly check that the headers have the extern "C" for C++
        # compilation.
        for I in sorted(rincludes - non_cxx_headers):
            with open(os.path.join(incdir,I)) as F:
                if 'extern "C" {' not in F.read():
                    raise ValueError("No extern C in %r"%(I));

        compile_test_headers(tmpd,incdir,includes,with_cxx=True);

# -------------------------------------------------------------------------


def get_symbol_names(fn):
    """Return the defined, public, symbols from a ELF shlib"""
    syms = subprocess.check_output(["readelf", "--wide", "-s", fn]).decode()
    go = False
    res = set()
    for I in syms.splitlines():
        if I.startswith("Symbol table '.dynsym'"):
            go = True
            continue

        if I.startswith(" ") and go:
            g = re.match(
                r"\s+\d+:\s+[0-9a-f]+\s+\d+.*(?:FUNC|OBJECT)\s+GLOBAL\s+DEFAULT\s+\d+\s+(\S+)@@(\S+)$",
                I)
            if not g or "PRIVATE" in g.group(2):
                continue
            res.add(g.group(1))
        else:
            go = False

    return res


def get_cc_args_from_pkgconfig(args, name, static):
    """Get the compile arguments from pkg-config for the named librarary"""
    os.environ["PKG_CONFIG_PATH"] = os.path.join(args.BUILD, "lib",
                                                 "pkgconfig")
    flags = ["pkg-config", "--errors-to-stdout", "--cflags", "--libs"]
    if static:
        flags.append("--static")
    opts = subprocess.check_output(flags + ["lib" + name]).decode()
    opts = shlex.split(opts)

    opts.insert(0, "-Wall")
    opts.insert(0, "-Werror")
    opts.insert(0, "-L%s" % (os.path.join(args.BUILD, "lib")))
    opts.insert(1, "-I%s" % (os.path.join(args.BUILD, "include")))
    if not static:
        return opts

    # Only static link the pkg-config stuff, otherwise we get warnings about
    # static linking portions of glibc that need NSS.
    opts.insert(0, "-Wl,-Bstatic")
    opts.append("-Wl,-Bdynamic")

    # We need this extra libpthread/m because libnl's pkgconfig file is
    # broken and doesn't include the private libraries it requires. :(
    if "-lnl-3" in opts:
        opts.append("-lm")
        opts.append("-lpthread")

    # Put glibc associated libraries out side the static link section,
    if "-lpthread" in opts:
        while "-lpthread" in opts:
            opts.remove("-lpthread")
        opts.append("-lpthread")
    if "-lm" in opts:
        while "-lm" in opts:
            opts.remove("-lm")
        opts.append("-lm")
    return opts


def compile_ninja(args, Fninja, name, cfn, opts):
    print("""
rule comp_{name}
    command = {CC} -Wall -o $out $in {opts}
    description = Compile and link $out
build {name} : comp_{name} {cfn}
default {name}""".format(
        name=name,
        CC=args.CC,
        cfn=cfn,
        opts=" ".join(pipes.quote(I) for I in opts)), file=Fninja)


def get_providers(args):
    """Return a list of provider names"""
    return set(
        I for I in os.listdir(os.path.join(args.SRC, "providers"))
        if not I.startswith("."))


def check_static_lib(args, tmpd, Fninja, static_lib, shared_lib, name):
    syms = get_symbol_names(shared_lib)
    if not syms:
        return

    cfn = os.path.join(tmpd, "%s-test.c" % (name))
    with open(cfn, "wt") as F:
        F.write("#include <stdio.h>\n")
        for I in syms:
            F.write("extern void %s(void);\n" % (I))
        F.write("int main(int argc,const char *argv[]) {\n")
        for I in syms:
            F.write('printf("%%p",&%s);\n' % (I))
        F.write("return 0; }\n")

    compile_ninja(args, Fninja, "%s-static-out" % (name), cfn,
                  get_cc_args_from_pkgconfig(args, name, static=True))
    compile_ninja(args, Fninja, "%s-shared-out" % (name), cfn,
                  get_cc_args_from_pkgconfig(args, name, static=False))


def check_static_providers(args, tmpd, Fninja):
    """Test that expected values for RDMA_STATIC_PROVIDERS are accepted and the
    link works"""
    cfn = os.path.join(tmpd, "provider-test.c")
    with open(cfn, "wt") as F:
        F.write("#include <infiniband/verbs.h>\n")
        F.write("int main(int argc,const char *argv[]) {\n")
        F.write('ibv_get_device_list(NULL);\n')
        F.write("return 0; }\n")

    opts = get_cc_args_from_pkgconfig(
        args, "ibverbs", static=True)

    providers = get_providers(args)
    for I in sorted(providers | {
            "none",
            "all",
    }):
        compile_ninja(args, Fninja, "providers-%s-static-out" % (I), cfn,
                      ["-DRDMA_STATIC_PROVIDERS=%s" % (I)] + opts)

    compile_ninja(
        args, Fninja, "providers-static-out", cfn,
        ["-DRDMA_STATIC_PROVIDERS=%s" % (",".join(providers))] + opts)


def test_static_libs(args):
    """Compile then link statically and dynamically a dummy program that touches
    every symbol in the libraries using pkgconfig output to guide the link
    options. This tests that pkgconfig is setup properly and that all the
    magic with incorporating the internal libraries for static linking has
    done its job."""
    libd = os.path.join(args.BUILD, "lib")
    success = True
    libs = []
    with inDirectory(libd):
        fns = set(fn for fn in os.listdir(".") if not os.path.islink(fn))
        for static_lib in fns:
            g = re.match(r"lib(.+)\.a$", static_lib)
            if g:
                for shared_lib in fns:
                    if re.match(r"lib%s.*\.so" % (g.group(1)), shared_lib):
                        libs.append((os.path.join(libd, static_lib),
                                     os.path.join(libd, shared_lib),
                                     g.group(1)))
                        break
                else:
                    raise ValueError(
                        "Failed to find matching shared library for %r" %
                        (static_lib))

    with private_tmp() as tmpd:
        with open(os.path.join(tmpd, "build.ninja"), "wt") as Fninja:
            for I in libs:
                check_static_lib(args, tmpd, Fninja, I[0], I[1], I[2])
            check_static_providers(args, tmpd, Fninja)
        subprocess.check_call(["ninja"], cwd=tmpd)


# -------------------------------------------------------------------------

parser = argparse.ArgumentParser(description='Run build time tests')
parser.add_argument("--build",default=os.getcwd(),dest="BUILD",
                    help="Build directory to inpsect");
parser.add_argument("--src",default=None,dest="SRC",
                    help="Top of the source tree");
parser.add_argument("--cc",default="cc",dest="CC",
                    help="C compiler to use");
parser.add_argument("--cxx",default="c++",dest="CXX",
                    help="C++ compiler to use");
args = parser.parse_args();

if args.SRC is None:
    args.SRC = get_src_dir();
args.SRC = os.path.abspath(args.SRC);
args.PACKAGE_VERSION = get_package_version(args);

funcs = globals();
for k,v in list(funcs.items()):
    if k.startswith("test_") and inspect.isfunction(v):
        v(args);