Kumarak/embed compile command

src/blight/actions/__init__.py

@@ -6,6 +6,7 @@
 from .cc_for_cxx import CCForCXX
 from .demo import Demo
 from .embed_bitcode import EmbedBitcode
+from .embed_commands import EmbedCommands  # noqa: F401
 from .find_inputs import FindInputs
 from .find_outputs import FindOutputs
 from .ignore_flags import IgnoreFlags
@@ -30,4 +31,5 @@
     "Lint",
     "Record",
     "SkipStrip",
+    "EmbedCommands",
 ]

src/blight/actions/embed_commands.py

@@ -0,0 +1,105 @@
+"""
+The `EmbedCommands` action.
+
+`EmbedCommands` embeds JSON compile commands, including environment variables,
+inside of a custom section of each built object file. These sections make it
+into the final binary.
+"""
+
+import hashlib
+import json
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any, Dict
+
+from blight.action import CompilerAction
+from blight.enums import Lang
+from blight.tool import CompilerTool
+from blight.util import flock_append
+
+
+def cc_as_string(tool_record: Dict) -> str:
+    return json.dumps(tool_record).replace('\\', '\\\\').replace('"', '\\"')
+
+
+def add_to_envp(envp: Dict, key: str, value: Any) -> None:
+    if isinstance(value, str):
+        envp[key] = value
+    elif isinstance(value, list):
+        envp[key] = [v for v in value]
+    else:
+        envp[key] = value
+
+VALID_EXTENSIONS = ("c", "cc", "cxx", "cpp", "c++")
+
+def is_input_source_code(tool: CompilerTool) -> bool:
+    for file_path in tool.inputs:
+        extension = str(file_path).lower().split(".")[-1]
+        if extension in VALID_EXTENSIONS:
+            return True
+    return False
+
+
+def cc_as_dict(tool: CompilerTool) -> Dict:
+    env: Dict[str, Any] = {}
+    tool_dict = tool.asdict()
+    old_env = tool_dict["env"]
+    for key, value in old_env.items():
+        if key == "PS1" or key == "LS_COLORS" or key.startswith("BLIGHT"):
+            continue
+        add_to_envp(env, key, value)
+
+    for key, value in old_env.items():
+        if key.startswith("BLIGHT_WRAPPED_"):
+            add_to_envp(env, key[15:], value)
+
+    return {
+        "cwd": tool_dict["cwd"],
+        "env": env,
+        "args": [v for v in tool.args],
+        "canonicalized_args": [v for v in tool.canonicalized_args],
+        "wrapped_tool": shutil.which(tool.wrapped_tool()),
+        "lang": tool.lang.name,
+    }
+
+
+_VARIABLE_TEMPLATE = """
+#if !defined(__linux__)
+__attribute__((section(\"__DATA,.trailofbits_cc\")))
+#elif !defined(__clang__)
+__attribute__((section(\".trailofbits_cc, \\"S\\", @note;\\n#\")))
+#else
+__attribute__((section(\".trailofbits_cc\")))
+#endif
+__attribute__((used))
+static const char cc_{}[] = \"{}\";
+"""
+
+
+class EmbedCommands(CompilerAction):
+    def __init__(self, config: dict[str, str]) -> None:
+        super().__init__(config)
+
+    def _get_header_file(self, cmd_hash: str) -> str:
+        f = tempfile.NamedTemporaryFile(suffix=".h", delete=False)
+        return f.name
+
+    def before_run(self, tool: CompilerTool) -> None:
+        if not is_input_source_code(tool):
+            return
+
+        cc_string = cc_as_string(cc_as_dict(tool)).strip()
+        cmd_hash = hashlib.sha256(cc_string.encode()).hexdigest()
+        header_file = self._get_header_file(cmd_hash)
+        with flock_append(Path(header_file)) as io:
+            variable = _VARIABLE_TEMPLATE.format(cmd_hash, cc_string)
+            print(variable, file=io)
+
+        tool.args += [
+            "-include",
+            header_file,
+            "-Wno-overlength-strings",
+            "-Wno-error",
+            "-Wno-unknown-escape-sequence",
+        ]

test/actions/test_embed_commands.py

@@ -0,0 +1,10 @@
+from blight.actions import EmbedCommands
+from blight.tool import CC
+
+
+def test_embed_bitcode():
+    embed_bitcode = EmbedCommands({})
+    cc1 = CC(["-o", "foo"])
+    embed_bitcode.before_run(cc1)
+    cc2 = CC(["foo.S"])
+    embed_bitcode.before_run(cc2)