Haskell: Patching libc with Cabal postBuild hook

mypackage.nix

# This isn't a full nix file, just a rough example on how to use the Setup.hs
{

  glibc_patched = pkgs.glibc.overrideAttrs (old: {
    patches = (old.patches or []) ++ [
      # glibc patch for malloc so that we get correct MallocInfo accountin in EKG.
      # That works only when this glibc is preloaded into a binrary,
      # see `glibc_patched_shell_exe_start_prefix`/`WITH_PATCHED_GLIBC` below.
      ./extra-nixpkgs/malloc_info-Fix-missing-accounting-of-top-chunk.patch
    ];
  });
  glibc_patched_interpreter = "${glibc_patched}/lib/ld-linux-x86-64.so.2";
  # Prepend this in front of a (shell, not execve()!) command to run it with our patched glibc.
  glibc_patched_shell_exe_start_prefix = "LD_PRELOAD=${glibc_patched}/lib/libc.so.6 ${glibc_patched_interpreter}";

  # When these are set, our Setup.hs `postBuild` hook automatically calls patchelf
  glibc_patched_cabalHookShellExports = ''
    export LIBC_TO_REPLACE_LIB_PATH="${pkgs.glibc}/lib"
    export PATCHED_LIBC_FULL_LIB_PATH="${glibc_patched}/lib"
    export PATCHED_LIBC_FULL_INTERPRETER_PATH="${glibc_patched_interpreter}"
  '';

  myhaskellpackage = mkDerivation rec {
    # ...
    shellHook = ''
      ${glibc_patched_cabalHookShellExports}
    '';

    preConfigure = ''
      ${glibc_patched_cabalHookShellExports}
    '';
  };

}

Setup.hs

{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wall #-}

import           Control.Concurrent.Async (forConcurrently_)
import           Data.Foldable (for_)
import           Data.List (intercalate)
import qualified Data.Map as Map
import           Data.Semigroup ((<>))
import qualified Data.Text as T
import           Distribution.Simple (defaultMainWithHooks, simpleUserHooks)
import           Distribution.Simple.Setup (buildVerbosity, fromFlagOrDefault)
import           Distribution.Simple.UserHooks (UserHooks(..))
import           Distribution.Simple.Utils (notice)
import           Distribution.Types.BuildInfo (BuildInfo)
import           Distribution.Types.ComponentLocalBuildInfo (ComponentLocalBuildInfo(..))
import           Distribution.Types.ComponentName (ComponentName, componentNameString)
import           Distribution.Types.LocalBuildInfo (LocalBuildInfo(..))
import           Distribution.Types.UnqualComponentName (unUnqualComponentName)
import           Distribution.Verbosity (Verbosity, normal)
import           System.Environment (lookupEnv)
import qualified System.FilePath as FP
import           System.Posix.Files.ByteString (touchFile)
import           System.Process (callProcess, readProcess)
import qualified Data.ByteString.Char8 as BSC


-- | Replaces libc using `patchelf` based on env vars:
-- > LIBC_TO_REPLACE_LIB_PATH
-- > PATCHED_LIBC_FULL_LIB_PATH
-- > PATCHED_LIBC_FULL_INTERPRETER_PATH
-- in all built executables.
patchelfLibc :: LocalBuildInfo -> Verbosity -> IO ()
patchelfLibc LocalBuildInfo{ componentNameMap, buildDir } verbosity = do
  mbToReplaceLibPath <- lookupEnv "LIBC_TO_REPLACE_LIB_PATH" -- should end with "/lib"
  for_ mbToReplaceLibPath $ \toReplaceLibPath -> do
    let replacementLibPathEnvVar = "PATCHED_LIBC_FULL_LIB_PATH"
    let replacementInterpreterPathEnvVar = "PATCHED_LIBC_FULL_INTERPRETER_PATH"
    replacementLibPath <- do
      mbPath <- lookupEnv replacementLibPathEnvVar
      maybe (fail $ "postBuild hook: missing env var " ++ replacementLibPathEnvVar) return mbPath

    replacementInterpreterPath <- do
      mbPath <- lookupEnv replacementInterpreterPathEnvVar
      maybe (fail $ "postBuild hook: missing env var " ++ replacementInterpreterPathEnvVar) return mbPath

    let patchComponent :: ComponentName -> IO ()
        patchComponent componentName = do
          let name =
                maybe
                  (error $ "patchelfLibc: component has no unqualified name: " ++ name)
                  unUnqualComponentName
                  (componentNameString componentName)
          notice verbosity $ "Patching libc of " <> name
          let exe = buildDir FP.</> name FP.</> name
          out <- readProcess "patchelf" ["--print-rpath", exe] ""
          let oldRpaths = [ T.unpack $ T.strip $ T.pack p | p <- FP.splitSearchPath out ]
          let newRpaths =
                [ if r == toReplaceLibPath then replacementLibPath else r
                | r <- oldRpaths
                ]
          let newRpath = intercalate [FP.searchPathSeparator] newRpaths
          callProcess
            "patchelf"
            [ "--set-interpreter", replacementInterpreterPath
            , "--set-rpath", newRpath, exe
            ]

    forConcurrently_ (concat $ Map.elems componentNameMap) $ \componentLocalBuildInfo -> do
      case componentLocalBuildInfo of
        -- Nothing to patch for libraries
        LibComponentLocalBuildInfo{} -> return ()
        FLibComponentLocalBuildInfo{} -> return ()
        -- Patch all executable things
        ExeComponentLocalBuildInfo{ componentLocalName = n } -> patchComponent n
        TestComponentLocalBuildInfo{ componentLocalName = n } -> patchComponent n
        BenchComponentLocalBuildInfo{ componentLocalName = n } -> patchComponent n


main :: IO ()
main = do
  defaultMainWithHooks $
    simpleUserHooks
      { postBuild = \_args buildFlags _packageDescription localBuildInfo -> do
          -- Insert our patched libc
          patchelfLibc localBuildInfo (fromFlagOrDefault normal (buildVerbosity buildFlags))
      }