From 852d059e2f5845f572b442f52364d273b4a736f6 Mon Sep 17 00:00:00 2001
From: sterni <sternenseemann@systemli.org>
Date: Tue, 14 Sep 2021 23:50:01 +0200
Subject: [PATCH] feat(nix/nint): accept attribute set with stdout, stderr and
 exit

This extends the calling convention for nint in a non-breaking way: If
the called script returns an attribute set instead of a string the
following is done:

* If the attributes `stdout` and/or `stderr` exist, their content (which
  must be a string currently) is written to the respective output.

* If the attribute `exit` exists, nint will exit with the given exit
  code. Must be a number that can be converted to an `i32`. If it's
  missing, nint will exit without indicating an error.

Change-Id: I209cf178fee3d970fdea3b26e4049e944af47457
Reviewed-on: https://cl.tvl.fyi/c/depot/+/3547
Tested-by: BuildkiteCI
Reviewed-by: tazjin <mail@tazj.in>
---
 nix/nint/README.md | 12 ++++++++++--
 nix/nint/nint.rs   | 48 +++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/nix/nint/README.md b/nix/nint/README.md
index ddd8045f7..369a82761 100644
--- a/nix/nint/README.md
+++ b/nix/nint/README.md
@@ -15,8 +15,16 @@ to the following calling convention:
     program name at `builtins.head argv`.
   * Extra arguments can be manually passed as described below.
 
-* The return value should always be a string (throwing is also okay)
-  which is printed to stdout by `nint`.
+* The return value must either be
+
+  * A string which is rendered to `stdout`.
+
+  * An attribute set with the following optional attributes:
+
+    * `stdout`: A string that's rendered to `stdout`
+    * `stderr`: A string that's rendered to `stderr`
+    * `exit`: A number which is used as an exit code.
+      If missing, nint always exits with 0 (or equivalent).
 
 ## Usage
 
diff --git a/nix/nint/nint.rs b/nix/nint/nint.rs
index 823de4865..1fa4dccb4 100644
--- a/nix/nint/nint.rs
+++ b/nix/nint/nint.rs
@@ -5,6 +5,7 @@ use std::ffi::OsString;
 use std::os::unix::ffi::{OsStringExt, OsStrExt};
 use std::io::{Error, ErrorKind, Write, stdout, stderr};
 use std::process::Command;
+use std::convert::{TryFrom};
 
 fn render_nix_string(s: &OsString) -> OsString {
     let mut rendered = Vec::new();
@@ -40,6 +41,26 @@ fn render_nix_list(arr: &[OsString]) -> OsString {
     OsString::from_vec(rendered)
 }
 
+/// Slightly overkill helper macro which takes a `Map<String, Value>` obtained
+/// from `Value::Object` and an output name (`stderr` or `stdout`) as an
+/// identifier. If a value exists for the given output in the object it gets
+/// written to the appropriate output.
+macro_rules! handle_set_output {
+    ($map_name:ident, $output_name:ident) => {
+        match $map_name.get(stringify!($output_name)) {
+            Some(Value::String(s)) =>
+                $output_name().write_all(s.as_bytes()),
+            Some(_) => Err(
+                Error::new(
+                    ErrorKind::Other,
+                    format!("Attribute {} must be a string!", stringify!($output_name)),
+                )
+            ),
+            None => Ok(()),
+        }
+    }
+}
+
 fn main() -> std::io::Result<()> {
     let mut nix_args = Vec::new();
 
@@ -90,6 +111,7 @@ fn main() -> std::io::Result<()> {
         nix_args.push(render_nix_list(&argv[..]));
 
         nix_args.push(OsString::from("--eval"));
+        nix_args.push(OsString::from("--strict"));
         nix_args.push(OsString::from("--json"));
 
         nix_args.push(argv[0].clone());
@@ -100,7 +122,31 @@ fn main() -> std::io::Result<()> {
 
         match serde_json::from_slice(&run.stdout[..]) {
             Ok(Value::String(s)) => stdout().write_all(s.as_bytes()),
-            Ok(_) => Err(Error::new(ErrorKind::Other, "output must be a string")),
+            Ok(Value::Object(m)) => {
+                handle_set_output!(m, stdout)?;
+                handle_set_output!(m, stderr)?;
+
+                match m.get("exit") {
+                    Some(Value::Number(n)) => {
+                        let code = n.as_i64().and_then(|v| i32::try_from(v).ok());
+
+                        match code {
+                            Some(i) => std::process::exit(i),
+                            None => Err(
+                                Error::new(
+                                    ErrorKind::Other,
+                                    "Attribute exit is not an i32"
+                                )
+                            ),
+                        }
+                    },
+                    Some(_) => Err(
+                        Error::new(ErrorKind::Other, "exit must be a number")
+                    ),
+                    None => Ok(()),
+                }
+            },
+            Ok(_) => Err(Error::new(ErrorKind::Other, "output must be a string or an object")),
             _ => {
                 stderr().write_all(&run.stderr[..]);
                 Err(Error::new(ErrorKind::Other, "internal nix error"))