From 894ddfc8b60ee615966da9fc9f97bd618e449415 Mon Sep 17 00:00:00 2001
From: dobosevych <dobosevych@users.noreply.github.com>
Date: Tue, 12 Apr 2022 15:52:41 +0300
Subject: [PATCH] Added possibility to serialize and deserialize binary
 messages in json (#1516)

* Added possibility to serialize and deserialize binary messages in json

* Flake8 fixed

* Hypothesis added to improve test range. Fixed issue b'\x80' serialization.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Added docstring

* Fixed pylint

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 kombu/utils/json.py       | 22 ++++++++++++++++++++--
 requirements/test.txt     |  1 +
 t/unit/utils/test_json.py | 12 ++++++++++++
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/kombu/utils/json.py b/kombu/utils/json.py
index cedaa793..6b8e4d46 100644
--- a/kombu/utils/json.py
+++ b/kombu/utils/json.py
@@ -1,5 +1,6 @@
 """JSON Serialization Utilities."""
 
+import base64
 import datetime
 import decimal
 import json as stdjson
@@ -55,6 +56,14 @@ class JSONEncoder(_encoder_cls):
                 return o.isoformat()
             elif isinstance(o, textual):
                 return text_t(o)
+            elif isinstance(o, bytes):
+                try:
+                    return {"bytes": o.decode("utf-8"), "__bytes__": True}
+                except UnicodeDecodeError:
+                    return {
+                        "bytes": base64.b64encode(o).decode("utf-8"),
+                        "__base64__": True,
+                    }
             return super().default(o)
 
 
@@ -69,7 +78,16 @@ def dumps(s, _dumps=json.dumps, cls=None, default_kwargs=None, **kwargs):
                   **dict(default_kwargs, **kwargs))
 
 
-def loads(s, _loads=json.loads, decode_bytes=True):
+def object_hook(dct):
+    """Hook function to perform custom deserialization."""
+    if "__bytes__" in dct:
+        return dct["bytes"].encode("utf-8")
+    if "__base64__" in dct:
+        return base64.b64decode(dct["bytes"].encode("utf-8"))
+    return dct
+
+
+def loads(s, _loads=json.loads, decode_bytes=True, object_hook=object_hook):
     """Deserialize json from string."""
     # None of the json implementations supports decoding from
     # a buffer/memoryview, or even reading from a stream
@@ -85,7 +103,7 @@ def loads(s, _loads=json.loads, decode_bytes=True):
         s = s.decode('utf-8')
 
     try:
-        return _loads(s)
+        return _loads(s, object_hook=object_hook)
     except _DecodeError:
         # catch "Unpaired high surrogate" error
         return stdjson.loads(s)
diff --git a/requirements/test.txt b/requirements/test.txt
index 9b8ca320..be41d4e5 100644
--- a/requirements/test.txt
+++ b/requirements/test.txt
@@ -2,3 +2,4 @@ pytz>dev
 pytest~=7.0.1
 pytest-sugar
 Pyro4
+hypothesis
diff --git a/t/unit/utils/test_json.py b/t/unit/utils/test_json.py
index 6af1c13b..d6a9c0b0 100644
--- a/t/unit/utils/test_json.py
+++ b/t/unit/utils/test_json.py
@@ -6,6 +6,8 @@ from uuid import uuid4
 
 import pytest
 import pytz
+from hypothesis import given, settings
+from hypothesis import strategies as st
 
 from kombu.utils.encoding import str_to_bytes
 from kombu.utils.json import _DecodeError, dumps, loads
@@ -39,6 +41,16 @@ class test_JSONEncoder:
             'date': stripped.isoformat(),
         }
 
+    @given(message=st.binary())
+    @settings(print_blob=True)
+    def test_binary(self, message):
+        serialized = loads(dumps({
+            'args': (message,),
+        }))
+        assert serialized == {
+            'args': [message],
+        }
+
     def test_Decimal(self):
         d = Decimal('3314132.13363235235324234123213213214134')
         assert loads(dumps({'d': d})), {'d': str(d)}