feat(Profpatsch): add netstring helper function & python lib r/2055

Simple helper function to generate a netstring that is a list of
key-value pairs, to serialize a nix dict. Also adds a python lib to
read the serialized form into a dict again.

Change-Id: I306c0cfd51640c0658d32c8d3a4f3d332ba448f0
Reviewed-on: https://cl.tvl.fyi/c/depot/+/2315
Tested-by: BuildkiteCI
Reviewed-by: Profpatsch <mail@profpatsch.de>

2 files changed, 103 insertions, 0 deletions

diff --git a/users/Profpatsch/netstring/default.nix b/users/Profpatsch/netstring/default.nix
new file mode 100644
index 0000000000..35345978a3
--- /dev/null
+++ b/users/Profpatsch/netstring/default.nix
@@ -0,0 +1,58 @@
+{ lib, pkgs, depot, ... }:
+let
+  toNetstring = s:
+    "${toString (builtins.stringLength s)}:${s},";
+
+  toNetstringKeyVal = attrs:
+    lib.concatStrings
+      (lib.mapAttrsToList
+        (k: v: toNetstring (toNetstring k + toNetstring v))
+        attrs);
+
+  python-netstring = depot.users.Profpatsch.writers.python3Lib {
+    name = "netstring";
+  } ''
+    def read_netstring(bytes):
+        (int_length, rest) = bytes.split(sep=b':', maxsplit=1)
+        val = rest[:int(int_length)]
+        # has to end on a ,
+        assert(rest[len(val)] == ord(','))
+        return (val, rest[len(val) + 1:])
+
+    def read_netstring_key_val(bytes):
+        (keyvalnet, rest) = read_netstring(bytes)
+        (key, valnet) = read_netstring(keyvalnet)
+        (val, nothing) = read_netstring(valnet)
+        assert(nothing == b"")
+        return (key, val, rest)
+
+    def read_netstring_key_val_list(bytes):
+        rest = bytes
+        res = {}
+        while rest != b"":
+            (key, val, r) = read_netstring_key_val(rest)
+            rest = r
+            res[key] = val
+        return res
+  '';
+
+  tests = import ./tests.nix {
+    inherit
+      depot
+      pkgs
+      lib
+      python-netstring
+      toNetstring
+      toNetstringKeyVal
+      ;
+  };
+
+in {
+  inherit
+    toNetstring
+    toNetstringKeyVal
+    python-netstring
+    tests
+      ;
+
+}
diff --git a/users/Profpatsch/netstring/tests.nix b/users/Profpatsch/netstring/tests.nix
new file mode 100644
index 0000000000..26853e13ee
--- /dev/null
+++ b/users/Profpatsch/netstring/tests.nix
@@ -0,0 +1,45 @@
+{ depot, lib, pkgs, python-netstring, toNetstring, toNetstringKeyVal }:
+
+let
+  imports = {
+    inherit (depot.users.Profpatsch)
+      writers
+      ;
+  };
+
+  python-netstring-test = imports.writers.python3 "python-netstring" {
+    libraries = p: [
+      python-netstring
+    ];
+  } ''
+    import netstring
+
+    def assEq(left, right):
+      assert left == right, "{} /= {}".format(str(left), str(right))
+
+    assEq(
+      netstring.read_netstring(b"""${toNetstring "hi!"}"""),
+      (b"hi!", b"")
+    )
+
+    assEq(
+      netstring.read_netstring_key_val(
+        b"""${toNetstringKeyVal { foo = "42"; }}"""
+      ),
+      (b'foo', b'42', b"")
+    )
+
+    assEq(
+      netstring.read_netstring_key_val_list(
+        b"""${toNetstringKeyVal { foo = "42"; bar = "hi"; }}"""
+      ),
+      { b'foo': b'42', b'bar': b'hi' }
+    )
+
+  '';
+
+in {
+  inherit
+    python-netstring-test
+    ;
+}