From 96231fab5f292ec9c58b22fd9026b43fde68adc7 Mon Sep 17 00:00:00 2001
From: Itamar Turner-Trauring <itamar@itamarst.org>
Date: Mon, 5 Oct 2020 11:01:11 -0400
Subject: [PATCH] Support bytes in JSON output.

---
 src/allmydata/test/test_util.py    | 19 +++++++++++
 src/allmydata/util/_python3.py     |  1 +
 src/allmydata/util/jsonbytes.py    | 51 ++++++++++++++++++++++++++++++
 src/allmydata/web/check_results.py |  2 +-
 4 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 src/allmydata/util/jsonbytes.py

diff --git a/src/allmydata/test/test_util.py b/src/allmydata/test/test_util.py
index f1f2b1c66..14fc020b9 100644
--- a/src/allmydata/test/test_util.py
+++ b/src/allmydata/test/test_util.py
@@ -13,16 +13,19 @@ if PY2:
 import six
 import os, time, sys
 import yaml
+import json
 
 from twisted.trial import unittest
 
 from allmydata.util import idlib, mathutil
 from allmydata.util import fileutil
+from allmydata.util import jsonbytes
 from allmydata.util import pollmixin
 from allmydata.util import yamlutil
 from allmydata.util.fileutil import EncryptedTemporaryFile
 from allmydata.test.common_util import ReallyEqualMixin
 
+
 if six.PY3:
     long = int
 
@@ -469,3 +472,19 @@ class YAML(unittest.TestCase):
         self.assertIsInstance(back[0], str)
         self.assertIsInstance(back[1], str)
         self.assertIsInstance(back[2], str)
+
+
+class JSONBytes(unittest.TestCase):
+    """Tests for BytesJSONEncoder."""
+
+    def test_encode_bytes(self):
+        """BytesJSONEncoder can encode bytes."""
+        data = {
+            b"hello": [1, b"cd"],
+        }
+        expected = {
+            u"hello": [1, u"cd"],
+        }
+        encoded = jsonbytes.dumps(data)
+        self.assertEqual(json.loads(encoded), expected)
+        self.assertEqual(jsonbytes.loads(encoded), expected)
diff --git a/src/allmydata/util/_python3.py b/src/allmydata/util/_python3.py
index ccff0958a..7ca18da15 100644
--- a/src/allmydata/util/_python3.py
+++ b/src/allmydata/util/_python3.py
@@ -72,6 +72,7 @@ PORTED_MODULES = [
     "allmydata.util.hashutil",
     "allmydata.util.humanreadable",
     "allmydata.util.iputil",
+    "allmydata.util.jsonbytes",
     "allmydata.util.log",
     "allmydata.util.mathutil",
     "allmydata.util.namespace",
diff --git a/src/allmydata/util/jsonbytes.py b/src/allmydata/util/jsonbytes.py
new file mode 100644
index 000000000..406a471a0
--- /dev/null
+++ b/src/allmydata/util/jsonbytes.py
@@ -0,0 +1,51 @@
+"""
+A JSON encoder than can serialize bytes.
+
+Ported to Python 3.
+"""
+
+from __future__ import unicode_literals
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from future.utils import PY2
+if PY2:
+    from future.builtins import filter, map, zip, ascii, chr, hex, input, next, oct, open, pow, round, super, bytes, dict, list, object, range, str, max, min  # noqa: F401
+
+
+import json
+
+
+class BytesJSONEncoder(json.JSONEncoder):
+    """
+    A JSON encoder than can also encode bytes.
+
+    The bytes are assumed to be UTF-8 encoded Unicode strings.
+    """
+    def default(self, o):
+        if isinstance(o, bytes):
+            return o.decode("utf-8")
+        return json.JSONEncoder.default(self, o)
+
+
+def dumps(obj, *args, **kwargs):
+    """Encode to JSON, supporting bytes as keys or values.
+
+    The bytes are assumed to be UTF-8 encoded Unicode strings.
+    """
+    if isinstance(obj, dict):
+        new_obj = {}
+        for k, v in obj.items():
+            if isinstance(k, bytes):
+                k = k.decode("utf-8")
+            new_obj[k] = v
+        obj = new_obj
+    return json.dumps(obj, cls=BytesJSONEncoder, *args, **kwargs)
+
+
+# To make this module drop-in compatible with json module:
+loads = json.loads
+
+
+__all__ = ["dumps", "loads"]
diff --git a/src/allmydata/web/check_results.py b/src/allmydata/web/check_results.py
index de2762df1..c87254f0d 100644
--- a/src/allmydata/web/check_results.py
+++ b/src/allmydata/web/check_results.py
@@ -1,7 +1,6 @@
 from future.builtins import str
 
 import time
-import json
 
 from twisted.web import (
     http,
@@ -32,6 +31,7 @@ from allmydata.interfaces import (
 from allmydata.util import (
     base32,
     dictutil,
+    jsonbytes as json,  # Supporting dumping bytes
 )