Increase coverage of spack.relocate (#16475)

- add docstrings and make parameter names consistent in `relocate.py`
- Make `replace_prefix_*` and other functions private (as they are implementation details)
- remove unused function _replace_prefix_nullterm()

- Add unit tests for `relocate.py` functions
  - add patchelf to Travis and use it during tests
  - add hello_world fixture with a compiled binary, so we can test relocation

2 files changed, 266 insertions, 120 deletions

diff --git a/lib/spack/spack/relocate.py b/lib/spack/spack/relocate.py
index 41faf273cd..94fd2b58bf 100644
--- a/lib/spack/spack/relocate.py
+++ b/lib/spack/spack/relocate.py
@@ -11,6 +11,7 @@ import llnl.util.lang
 import llnl.util.tty as tty
 import macholib.MachO
 import macholib.mach_o
+import spack.architecture
 import spack.cmd
 import spack.repo
 import spack.spec
@@ -385,57 +386,79 @@ def macholib_get_paths(cur_path):
     return (rpaths, deps, ident)
 
 
-def modify_elf_object(path_name, new_rpaths):
-    """
-    Replace orig_rpath with new_rpath in RPATH of elf object path_name
-    """
+def _set_elf_rpaths(target, rpaths):
+    """Replace the original RPATH of the target with the paths passed
+    as arguments.
 
-    new_joined = ':'.join(new_rpaths)
+    This function uses ``patchelf`` to set RPATHs.
 
-    # if we're relocating patchelf itself, use it
-    bak_path = path_name + ".bak"
+    Args:
+        target: target executable. Must be an ELF object.
+        rpaths: paths to be set in the RPATH
 
-    if path_name[-13:] == "/bin/patchelf":
-        shutil.copy(path_name, bak_path)
-        patchelf = executable.Executable(bak_path)
-    else:
-        patchelf = executable.Executable(_patchelf())
+    Returns:
+        A string concatenating the stdout and stderr of the call
+        to ``patchelf``
+    """
+    # Join the paths using ':' as a separator
+    rpaths_str = ':'.join(rpaths)
+
+    # If we're relocating patchelf itself, make a copy and use it
+    bak_path = None
+    if target.endswith("/bin/patchelf"):
+        bak_path = target + ".bak"
+        shutil.copy(target, bak_path)
 
+    patchelf, output = executable.Executable(bak_path or _patchelf()), None
     try:
-        patchelf('--force-rpath', '--set-rpath', '%s' % new_joined,
-                 '%s' % path_name, output=str, error=str)
+        # TODO: revisit the use of --force-rpath as it might be conditional
+        # TODO: if we want to support setting RUNPATH from binary packages
+        patchelf_args = ['--force-rpath', '--set-rpath', rpaths_str, target]
+        output = patchelf(*patchelf_args, output=str, error=str)
     except executable.ProcessError as e:
-        msg = 'patchelf --force-rpath --set-rpath %s failed with error %s' % (
-            path_name, e)
-        tty.warn(msg)
-    if os.path.exists(bak_path):
-        os.remove(bak_path)
+        msg = 'patchelf --force-rpath --set-rpath {0} failed with error {1}'
+        tty.warn(msg.format(target, e))
+    finally:
+        if bak_path and os.path.exists(bak_path):
+            os.remove(bak_path)
+    return output
 
 
 def needs_binary_relocation(m_type, m_subtype):
-    """
-    Check whether the given filetype is a binary that may need relocation.
+    """Returns True if the file with MIME type/subtype passed as arguments
+    needs binary relocation, False otherwise.
+
+    Args:
+        m_type (str): MIME type of the file
+        m_subtype (str): MIME subtype of the file
     """
     if m_type == 'application':
-        if (m_subtype == 'x-executable' or m_subtype == 'x-sharedlib' or
-                m_subtype == 'x-mach-binary'):
+        if m_subtype in ('x-executable', 'x-sharedlib', 'x-mach-binary'):
             return True
     return False
 
 
 def needs_text_relocation(m_type, m_subtype):
+    """Returns True if the file with MIME type/subtype passed as arguments
+    needs text relocation, False otherwise.
+
+    Args:
+        m_type (str): MIME type of the file
+        m_subtype (str): MIME subtype of the file
     """
-    Check whether the given filetype is text that may need relocation.
-    """
-    return (m_type == "text")
+    return m_type == 'text'
 
 
-def replace_prefix_text(path_name, old_dir, new_dir):
-    """
-    Replace old install prefix with new install prefix
-    in text files using utf-8 encoded strings.
+def _replace_prefix_text(filename, old_dir, new_dir):
+    """Replace all the occurrences of the old install prefix with a
+    new install prefix in text files that are utf-8 encoded.
+
+    Args:
+        filename (str): target text file (utf-8 encoded)
+        old_dir (str): directory to be searched in the file
+        new_dir (str): substitute for the old directory
     """
-    with open(path_name, 'rb+') as f:
+    with open(filename, 'rb+') as f:
         data = f.read()
         f.seek(0)
         # Replace old_dir with new_dir if it appears at the beginning of a path
@@ -454,47 +477,18 @@ def replace_prefix_text(path_name, old_dir, new_dir):
         f.truncate()
 
 
-def replace_prefix_bin(path_name, old_dir, new_dir):
-    """
-    Attempt to replace old install prefix with new install prefix
-    in binary files by prefixing new install prefix with os.sep
-    until the lengths of the prefixes are the same.
-    """
+def _replace_prefix_bin(filename, old_dir, new_dir):
+    """Replace all the occurrences of the old install prefix with a
+    new install prefix in binary files.
 
-    def replace(match):
-        occurances = match.group().count(old_dir.encode('utf-8'))
-        olen = len(old_dir.encode('utf-8'))
-        nlen = len(new_dir.encode('utf-8'))
-        padding = (olen - nlen) * occurances
-        if padding < 0:
-            return data
-        return match.group().replace(old_dir.encode('utf-8'),
-                                     os.sep.encode('utf-8') * padding +
-                                     new_dir.encode('utf-8'))
+    The new install prefix is prefixed with ``os.sep`` until the
+    lengths of the prefixes are the same.
 
-    with open(path_name, 'rb+') as f:
-        data = f.read()
-        f.seek(0)
-        original_data_len = len(data)
-        pat = re.compile(old_dir.encode('utf-8'))
-        if not pat.search(data):
-            return
-        ndata = pat.sub(replace, data)
-        if not len(ndata) == original_data_len:
-            raise BinaryStringReplacementError(
-                path_name, original_data_len, len(ndata))
-        f.write(ndata)
-        f.truncate()
-
-
-def replace_prefix_nullterm(path_name, old_dir, new_dir):
-    """
-    Attempt to replace old install prefix with new install prefix
-    in binary files by replacing with null terminated string
-    that is the same length unless the old path is shorter
-    Used on linux to replace mach-o rpaths
+    Args:
+        filename (str): target binary file
+        old_dir (str): directory to be searched in the file
+        new_dir (str): substitute for the old directory
     """
-
     def replace(match):
         occurances = match.group().count(old_dir.encode('utf-8'))
         olen = len(old_dir.encode('utf-8'))
@@ -502,23 +496,22 @@ def replace_prefix_nullterm(path_name, old_dir, new_dir):
         padding = (olen - nlen) * occurances
         if padding < 0:
             return data
-        return match.group().replace(old_dir.encode('utf-8'),
-                                     new_dir.encode('utf-8')) + b'\0' * padding
+        return match.group().replace(
+            old_dir.encode('utf-8'),
+            os.sep.encode('utf-8') * padding + new_dir.encode('utf-8')
+        )
 
-    if len(new_dir) > len(old_dir):
-        raise BinaryTextReplaceError(old_dir, new_dir)
-
-    with open(path_name, 'rb+') as f:
+    with open(filename, 'rb+') as f:
         data = f.read()
         f.seek(0)
         original_data_len = len(data)
-        pat = re.compile(re.escape(old_dir).encode('utf-8') + b'([^\0]*?)\0')
+        pat = re.compile(old_dir.encode('utf-8'))
         if not pat.search(data):
             return
         ndata = pat.sub(replace, data)
         if not len(ndata) == original_data_len:
             raise BinaryStringReplacementError(
-                path_name, original_data_len, len(ndata))
+                filename, original_data_len, len(ndata))
         f.write(ndata)
         f.truncate()
 
@@ -597,50 +590,89 @@ def relocate_macho_binaries(path_names, old_layout_root, new_layout_root,
                                        paths_to_paths)
 
 
-def elf_find_paths(orig_rpaths, old_layout_root, prefix_to_prefix):
-    new_rpaths = list()
+def _transform_rpaths(orig_rpaths, orig_root, new_prefixes):
+    """Return an updated list of RPATHs where each entry in the original list
+    starting with the old root is relocated to another place according to the
+    mapping passed as argument.
+
+    Args:
+        orig_rpaths (list): list of the original RPATHs
+        orig_root (str): original root to be substituted
+        new_prefixes (dict): dictionary that maps the original prefixes to
+            where they should be relocated
+
+    Returns:
+        List of paths
+    """
+    new_rpaths = []
     for orig_rpath in orig_rpaths:
-        if orig_rpath.startswith(old_layout_root):
-            for old_prefix, new_prefix in prefix_to_prefix.items():
-                if orig_rpath.startswith(old_prefix):
-                    new_rpaths.append(re.sub(re.escape(old_prefix),
-                                             new_prefix, orig_rpath))
-        else:
+        # If the original RPATH doesn't start with the target root
+        # append it verbatim and proceed
+        if not orig_rpath.startswith(orig_root):
             new_rpaths.append(orig_rpath)
+            continue
+
+        # Otherwise inspect the mapping and transform + append any prefix
+        # that starts with a registered key
+        for old_prefix, new_prefix in new_prefixes.items():
+            if orig_rpath.startswith(old_prefix):
+                new_rpaths.append(
+                    re.sub(re.escape(old_prefix), new_prefix, orig_rpath)
+                )
+
     return new_rpaths
 
 
-def relocate_elf_binaries(path_names, old_layout_root, new_layout_root,
-                          prefix_to_prefix, rel, old_prefix, new_prefix):
-    """
-    Use patchelf to get the original rpaths and then replace them with
+def relocate_elf_binaries(binaries, orig_root, new_root,
+                          new_prefixes, rel, orig_prefix, new_prefix):
+    """Relocate the binaries passed as arguments by changing their RPATHs.
+
+    Use patchelf to get the original RPATHs and then replace them with
     rpaths in the new directory layout.
-    New rpaths are determined from a dictionary mapping the prefixes in the
+
+    New RPATHs are determined from a dictionary mapping the prefixes in the
     old directory layout to the prefixes in the new directory layout if the
     rpath was in the old layout root, i.e. system paths are not replaced.
-    """
-    for path_name in path_names:
-        orig_rpaths = _elf_rpaths_for(path_name)
-        new_rpaths = list()
+
+    Args:
+        binaries (list): list of binaries that might need relocation, located
+            in the new prefix
+        orig_root (str): original root to be substituted
+        new_root (str): new root to be used, only relevant for relative RPATHs
+        new_prefixes (dict): dictionary that maps the original prefixes to
+            where they should be relocated
+        rel (bool): True if the RPATHs are relative, False if they are absolute
+        orig_prefix (str): prefix where the executable was originally located
+        new_prefix (str): prefix where we want to relocate the executable
+    """
+    for new_binary in binaries:
+        orig_rpaths = _elf_rpaths_for(new_binary)
+        # TODO: Can we deduce `rel` from the original RPATHs?
         if rel:
-            # get the file path in the old_prefix
-            orig_path_name = re.sub(re.escape(new_prefix), old_prefix,
-                                    path_name)
-            # get the normalized rpaths in the old prefix using the file path
+            # Get the file path in the original prefix
+            orig_binary = re.sub(
+                re.escape(new_prefix), orig_prefix, new_binary
+            )
+
+            # Get the normalized RPATHs in the old prefix using the file path
             # in the orig prefix
-            orig_norm_rpaths = _normalize_relative_paths(orig_path_name,
-                                                         orig_rpaths)
-            # get the normalize rpaths in the new prefix
-            norm_rpaths = elf_find_paths(orig_norm_rpaths, old_layout_root,
-                                         prefix_to_prefix)
-            # get the relativized rpaths in the new prefix
-            new_rpaths = _make_relative(path_name, new_layout_root,
-                                        norm_rpaths)
-            modify_elf_object(path_name, new_rpaths)
+            orig_norm_rpaths = _normalize_relative_paths(
+                orig_binary, orig_rpaths
+            )
+            # Get the normalize RPATHs in the new prefix
+            new_norm_rpaths = _transform_rpaths(
+                orig_norm_rpaths, orig_root, new_prefixes
+            )
+            # Get the relative RPATHs in the new prefix
+            new_rpaths = _make_relative(
+                new_binary, new_root, new_norm_rpaths
+            )
+            _set_elf_rpaths(new_binary, new_rpaths)
         else:
-            new_rpaths = elf_find_paths(orig_rpaths, old_layout_root,
-                                        prefix_to_prefix)
-            modify_elf_object(path_name, new_rpaths)
+            new_rpaths = _transform_rpaths(
+                orig_rpaths, orig_root, new_prefixes
+            )
+            _set_elf_rpaths(new_binary, new_rpaths)
 
 
 def make_link_relative(cur_path_names, orig_path_names):
@@ -684,7 +716,7 @@ def make_elf_binaries_relative(cur_path_names, orig_path_names,
         if orig_rpaths:
             new_rpaths = _make_relative(orig_path, old_layout_root,
                                         orig_rpaths)
-            modify_elf_object(cur_path, new_rpaths)
+            _set_elf_rpaths(cur_path, new_rpaths)
 
 
 def check_files_relocatable(cur_path_names, allow_root):
@@ -738,11 +770,11 @@ def relocate_text(path_names, old_layout_root, new_layout_root,
     sbangnew = '#!/bin/bash %s/bin/sbang' % new_spack_prefix
 
     for path_name in path_names:
-        replace_prefix_text(path_name, old_install_prefix, new_install_prefix)
+        _replace_prefix_text(path_name, old_install_prefix, new_install_prefix)
         for orig_dep_prefix, new_dep_prefix in prefix_to_prefix.items():
-            replace_prefix_text(path_name, orig_dep_prefix, new_dep_prefix)
-        replace_prefix_text(path_name, old_layout_root, new_layout_root)
-        replace_prefix_text(path_name, sbangre, sbangnew)
+            _replace_prefix_text(path_name, orig_dep_prefix, new_dep_prefix)
+        _replace_prefix_text(path_name, old_layout_root, new_layout_root)
+        _replace_prefix_text(path_name, sbangre, sbangnew)
 
 
 def relocate_text_bin(path_names, old_layout_root, new_layout_root,
@@ -758,9 +790,9 @@ def relocate_text_bin(path_names, old_layout_root, new_layout_root,
         for path_name in path_names:
             for old_dep_prefix, new_dep_prefix in prefix_to_prefix.items():
                 if len(new_dep_prefix) <= len(old_dep_prefix):
-                    replace_prefix_bin(
+                    _replace_prefix_bin(
                         path_name, old_dep_prefix, new_dep_prefix)
-            replace_prefix_bin(path_name, old_spack_prefix, new_spack_prefix)
+            _replace_prefix_bin(path_name, old_spack_prefix, new_spack_prefix)
     else:
         if len(path_names) > 0:
             raise BinaryTextReplaceError(
diff --git a/lib/spack/spack/test/relocate.py b/lib/spack/spack/test/relocate.py
index a2bba52f40..a594f1a9ba 100644
--- a/lib/spack/spack/test/relocate.py
+++ b/lib/spack/spack/test/relocate.py
@@ -103,6 +103,31 @@ def mock_patchelf(tmpdir):
     return _factory
 
 
+@pytest.fixture()
+def hello_world(tmpdir):
+    source = tmpdir.join('main.c')
+    source.write("""
+#include <stdio.h>
+int main(){
+    printf("Hello world!");
+}
+""")
+
+    def _factory(rpaths):
+        gcc = spack.util.executable.which('gcc')
+        executable = source.dirpath('main.x')
+        rpath_str = ':'.join(rpaths)
+        opts = [
+            '-Wl,--disable-new-dtags',
+            '-Wl,-rpath={0}'.format(rpath_str),
+            str(source), '-o', str(executable)
+        ]
+        gcc(*opts)
+        return executable
+
+    return _factory
+
+
 @pytest.mark.requires_executables(
     '/usr/bin/gcc', 'patchelf', 'strings', 'file'
 )
@@ -118,9 +143,7 @@ def test_file_is_relocatable(source_file, is_relocatable):
     assert spack.relocate.file_is_relocatable(executable) is is_relocatable
 
 
-@pytest.mark.requires_executables(
-    'patchelf', 'strings', 'file'
-)
+@pytest.mark.requires_executables('patchelf', 'strings', 'file')
 def test_patchelf_is_relocatable():
     patchelf = spack.relocate._patchelf()
     assert llnl.util.filesystem.is_exe(patchelf)
@@ -194,3 +217,94 @@ def test_normalize_relative_paths(start_path, relative_paths, expected):
         start_path, relative_paths
     )
     assert normalized == expected
+
+
+def test_set_elf_rpaths(mock_patchelf):
+    # Try to relocate a mock version of patchelf and check
+    # the call made to patchelf itself
+    patchelf = mock_patchelf('echo $@')
+    rpaths = ['/usr/lib', '/usr/lib64', '/opt/local/lib']
+    output = spack.relocate._set_elf_rpaths(patchelf, rpaths)
+
+    # Assert that the arguments of the call to patchelf are as expected
+    assert '--force-rpath' in output
+    assert '--set-rpath ' + ':'.join(rpaths) in output
+    assert patchelf in output
+
+
+def test_set_elf_rpaths_warning(mock_patchelf):
+    # Mock a failing patchelf command and ensure it warns users
+    patchelf = mock_patchelf('exit 1')
+    rpaths = ['/usr/lib', '/usr/lib64', '/opt/local/lib']
+    # To avoid using capfd in order to check if the warning was triggered
+    # here we just check that output is not set
+    output = spack.relocate._set_elf_rpaths(patchelf, rpaths)
+    assert output is None
+
+
+@pytest.mark.requires_executables('patchelf', 'strings', 'file', 'gcc')
+def test_replace_prefix_bin(hello_world):
+    # Compile an "Hello world!" executable and set RPATHs
+    executable = hello_world(rpaths=['/usr/lib', '/usr/lib64'])
+
+    # Relocate the RPATHs
+    spack.relocate._replace_prefix_bin(str(executable), '/usr', '/foo')
+
+    # Check that the RPATHs changed
+    patchelf = spack.util.executable.which('patchelf')
+    output = patchelf('--print-rpath', str(executable), output=str)
+    assert output.strip() == '/foo/lib:/foo/lib64'
+
+
+@pytest.mark.requires_executables('patchelf', 'strings', 'file', 'gcc')
+def test_relocate_elf_binaries_absolute_paths(hello_world, tmpdir):
+    # Create an executable, set some RPATHs, copy it to another location
+    orig_binary = hello_world(rpaths=[str(tmpdir.mkdir('lib')), '/usr/lib64'])
+    new_root = tmpdir.mkdir('another_dir')
+    shutil.copy(str(orig_binary), str(new_root))
+
+    # Relocate the binary
+    new_binary = new_root.join('main.x')
+    spack.relocate.relocate_elf_binaries(
+        binaries=[str(new_binary)],
+        orig_root=str(orig_binary.dirpath()),
+        new_root=None,  # Not needed when relocating absolute paths
+        new_prefixes={
+            str(tmpdir): '/foo'
+        },
+        rel=False,
+        # Not needed when relocating absolute paths
+        orig_prefix=None, new_prefix=None
+    )
+
+    # Check that the RPATHs changed
+    patchelf = spack.util.executable.which('patchelf')
+    output = patchelf('--print-rpath', str(new_binary), output=str)
+    assert output.strip() == '/foo/lib:/usr/lib64'
+
+
+@pytest.mark.requires_executables('patchelf', 'strings', 'file', 'gcc')
+def test_relocate_elf_binaries_relative_paths(hello_world, tmpdir):
+    # Create an executable, set some RPATHs, copy it to another location
+    orig_binary = hello_world(
+        rpaths=['$ORIGIN/lib', '$ORIGIN/lib64', '/opt/local/lib']
+    )
+    new_root = tmpdir.mkdir('another_dir')
+    shutil.copy(str(orig_binary), str(new_root))
+
+    # Relocate the binary
+    new_binary = new_root.join('main.x')
+    spack.relocate.relocate_elf_binaries(
+        binaries=[str(new_binary)],
+        orig_root=str(orig_binary.dirpath()),
+        new_root=str(new_root),
+        new_prefixes={str(tmpdir): '/foo'},
+        rel=True,
+        orig_prefix=str(orig_binary.dirpath()),
+        new_prefix=str(new_root)
+    )
+
+    # Check that the RPATHs changed
+    patchelf = spack.util.executable.which('patchelf')
+    output = patchelf('--print-rpath', str(new_binary), output=str)
+    assert output.strip() == '/foo/lib:/foo/lib64:/opt/local/lib'