From b3e3bf9b931577adf067c1ae40d6ab7237249ed4 Mon Sep 17 00:00:00 2001
From: matthewcroughan <matt@croughan.sh>
Date: Sun, 8 Sep 2024 17:50:16 +0100
Subject: [PATCH] trivial-builders: init addTmateBreakpoint

---
 .../trivial-builders/default.nix              | 41 ++++++++++++++++++-
 .../test/addTmateBreakpoint.nix               |  1 +
 .../trivial-builders/test/default.nix         |  3 ++
 pkgs/top-level/stage.nix                      |  2 +-
 4 files changed, 45 insertions(+), 2 deletions(-)
 create mode 100644 pkgs/build-support/trivial-builders/test/addTmateBreakpoint.nix

diff --git a/pkgs/build-support/trivial-builders/default.nix b/pkgs/build-support/trivial-builders/default.nix
index c5b27e85e5bbac..00612260f18003 100644
--- a/pkgs/build-support/trivial-builders/default.nix
+++ b/pkgs/build-support/trivial-builders/default.nix
@@ -1,4 +1,4 @@
-{ lib, config, stdenv, stdenvNoCC, jq, lndir, runtimeShell, shellcheck-minimal }:
+{ lib, config, stdenv, stdenvNoCC, jq, lndir, unixtools, tmate, runtimeShell, shellcheck-minimal }:
 
 let
   inherit (lib)
@@ -602,6 +602,45 @@ rec {
     ln -s ${lib.getBin drv}/bin $out/bin
   '';
 
+  /*
+    Adds a hook to the nativeBuildInputs of a derivation which will provide an
+    SSH session via tmate to the build environment, for interactive
+    debugging
+
+    Example:
+
+    # Spawn a tmate session after the `buildPhase` of the `hello` derivation
+    addTmateBreakpoint (pkgs.hello.overrideAttrs { buildPhase = "exit 1;"; })
+  */
+  addTmateBreakpoint = let
+    scriptFlags = if stdenv.hostPlatform.isDarwin then "-qeF ./tmpfile" else "-qefc";
+    warningMessage = "addTmateBreakpoint in use, ignore the sha256 warning and do not leak these build logs, otherwise unauthorized ssh access to the build sandbox may occur";
+    hook = lib.warn warningMessage (makeSetupHook {} (writeTextFile {
+      name = "hook.sh";
+      text = ''
+        breakpointHookTmate() {
+          local red='\033[0;31m'
+          local no_color='\033[0m'
+
+          echo -e "''${red}build failed in ''${curPhase} with exit code ''${exitCode}''${no_color}"
+          echo -e "''${red}### WARNING ###''${no_color}"
+          echo -e "''${red}${warningMessage}''${no_color}"
+
+          printf "To attach using tmate:\n\n"
+          ${unixtools.script}/bin/script ${scriptFlags} "${tmate}/bin/tmate -F"
+        }
+        failureHooks+=(breakpointHookTmate)
+      '';
+    }));
+  in drv: drv.overrideAttrs (old: {
+    name = "${drv.name}-addTmateBreakpoint-${builtins.unsafeDiscardStringContext (lib.substring 0 12 (baseNameOf drv.drvPath))}";
+    nativeBuildInputs = (old.nativeBuildInputs or []) ++ [
+      hook
+    ];
+    outputHash = "";
+    outputHashAlgo = "sha256";
+    outputHashMode = "recursive";
+  });
 
   # Docs in doc/build-helpers/special/makesetuphook.section.md
   # See https://nixos.org/manual/nixpkgs/unstable/#sec-pkgs.makeSetupHook
diff --git a/pkgs/build-support/trivial-builders/test/addTmateBreakpoint.nix b/pkgs/build-support/trivial-builders/test/addTmateBreakpoint.nix
new file mode 100644
index 00000000000000..b78088f4163263
--- /dev/null
+++ b/pkgs/build-support/trivial-builders/test/addTmateBreakpoint.nix
@@ -0,0 +1 @@
+{ addTmateBreakpoint, hello }: (addTmateBreakpoint hello).overrideAttrs { postPatch = "exit 1"; }
diff --git a/pkgs/build-support/trivial-builders/test/default.nix b/pkgs/build-support/trivial-builders/test/default.nix
index e1ed0be72bf358..0c48ae9406eece 100644
--- a/pkgs/build-support/trivial-builders/test/default.nix
+++ b/pkgs/build-support/trivial-builders/test/default.nix
@@ -33,4 +33,7 @@ recurseIntoAttrs {
     inherit (references) samples;
   };
   writeTextFile = callPackage ./write-text-file.nix {};
+  manualTesting = lib.dontRecurseIntoAttrs {
+    addTmateBreakpoint = callPackage ./addTmateBreakpoint.nix {};
+  };
 }
diff --git a/pkgs/top-level/stage.nix b/pkgs/top-level/stage.nix
index 265ab242d86d21..8266669c692498 100644
--- a/pkgs/top-level/stage.nix
+++ b/pkgs/top-level/stage.nix
@@ -114,7 +114,7 @@ let
       inherit lib;
       inherit (self) config;
       inherit (self) runtimeShell stdenv stdenvNoCC;
-      inherit (self.pkgsBuildHost) jq shellcheck-minimal;
+      inherit (self.pkgsBuildHost) jq shellcheck-minimal unixtools tmate;
       inherit (self.pkgsBuildHost.xorg) lndir;
     };