fix: setting bytes field from python string base64 decodes before assignment (#255)

This is necessary for maintaining field integrity during round trip
conversions to python dicts.

E.g.

    class MyMessage(proto.Message):
          data = proto.Field(proto.BYTES, number=1)

    my_message = MyMessage(data=b"this is a data payload")
    assert my_message == MyMessage(MyMessage.to_dict(my_message))

Conversion to and from json does not exhibit this problem because the
entire ser/des logic is encapuslated by the vanilla protobuf runtime,
which handles this problem.

Fix for #249

Author: software-dov

docs/marshal.rst

@@ -44,6 +44,34 @@ Protocol buffer type                Python type             Nullable
     If you *write* a timestamp field using a Python ``datetime`` value,
     any existing nanosecond precision will be overwritten.
 
+.. note::
+
+   Setting a ``bytes`` field from a string value will first base64 decode the string.
+   This is necessary to preserve the original protobuf semantics when converting between
+   Python dicts and proto messages.
+   Converting a message containing a bytes field to a dict will
+   base64 encode the bytes field and yield a value of type str.
+
+.. code-block:: python
+
+  import proto
+  from google.protobuf.json_format import ParseDict
+
+  class MyMessage(proto.Message):
+      data = proto.Field(proto.BYTES, number=1)
+
+  msg = MyMessage(data=b"this is a message")
+  msg_dict = MyMessage.to_dict(msg)
+
+  # Note: the value is the base64 encoded string of the bytes field.
+  # It has a type of str, NOT bytes.
+  assert type(msg_dict['data']) == str
+
+  msg_pb = ParseDict(msg_dict, MyMessage.pb())
+  msg_two = MyMessage(msg_dict)
+
+  assert msg == msg_pb == msg_two
+
 
 Wrapper types
 -------------

proto/fields.py

@@ -126,14 +126,15 @@ def package(self) -> str:
 
     @property
     def pb_type(self):
-        """Return the composite type of the field, or None for primitives."""
+        """Return the composite type of the field, or the primitive type if a primitive."""
         # For enums, return the Python enum.
         if self.enum:
             return self.enum
 
-        # For non-enum primitives, return None.
+        # For primitive fields, we still want to know
+        # what the type is.
         if not self.message:
-            return None
+            return self.proto_type
 
         # Return the internal protobuf message.
         if hasattr(self.message, "_meta"):

proto/marshal/marshal.py

@@ -25,9 +25,11 @@
 from proto.marshal.collections import MapComposite
 from proto.marshal.collections import Repeated
 from proto.marshal.collections import RepeatedComposite
+from proto.marshal.rules import bytes as pb_bytes
 from proto.marshal.rules import dates
 from proto.marshal.rules import struct
 from proto.marshal.rules import wrappers
+from proto.primitives import ProtoType
 
 
 class Rule(abc.ABC):
@@ -85,14 +87,6 @@ class TimestampRule:
             proto_type (type): A protocol buffer message type.
             rule: A marshal object
         """
-        # Sanity check: Do not register anything to a class that is not
-        # a protocol buffer message.
-        if not issubclass(proto_type, (message.Message, enum.IntEnum)):
-            raise TypeError(
-                "Only enums and protocol buffer messages may be "
-                "registered to the marshal."
-            )
-
         # If a rule was provided, register it and be done.
         if rule:
             # Ensure the rule implements Rule.
@@ -150,6 +144,9 @@ def reset(self):
         self.register(struct_pb2.ListValue, struct.ListValueRule(marshal=self))
         self.register(struct_pb2.Struct, struct.StructRule(marshal=self))
 
+        # Special case for bytes to allow base64 encode/decode
+        self.register(ProtoType.BYTES, pb_bytes.BytesRule())
+
     def to_python(self, proto_type, value, *, absent: bool = None):
         # Internal protobuf has its own special type for lists of values.
         # Return a view around it that implements MutableSequence.

proto/marshal/rules/bytes.py

@@ -0,0 +1,44 @@
+# Copyright (C) 2021  Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import base64
+
+
+class BytesRule:
+    """A marshal between Python strings and protobuf bytes.
+
+    Note: this conversion is asymmetric becasue Python does have a bytes type.
+    It is sometimes necessary to convert proto bytes fields to strings, e.g. for
+    JSON encoding, marshalling a message to a dict. Because bytes fields can
+    represent arbitrary data, bytes fields are base64 encoded when they need to
+    be represented as strings.
+
+    It is necessary to have the conversion be bidirectional, i.e.
+    my_message == MyMessage(MyMessage.to_dict(my_message))
+
+    To accomplish this, we need to intercept assignments from strings and
+    base64 decode them back into bytes.
+    """
+
+    def to_python(self, value, *, absent: bool = None):
+        return value
+
+    def to_proto(self, value):
+        if isinstance(value, str):
+            value = value.encode("utf-8")
+            value += b"=" * (4 - len(value) % 4)  # padding
+            value = base64.urlsafe_b64decode(value)
+
+        return value

tests/test_fields_bytes.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import base64
 import pytest
 
 import proto
@@ -71,5 +72,25 @@ class Foo(proto.Message):
     # for strings (but not vice versa).
     foo.bar = b"anything"
     assert foo.bar == "anything"
-    with pytest.raises(TypeError):
-        foo.baz = "anything"
+
+    # We need to permit setting bytes fields from strings,
+    # but the marshalling needs to base64 decode the result.
+    # This is a requirement for interop with the vanilla protobuf runtime:
+    # converting a proto message to a dict base64 encodes the bytes
+    # becase it may be sent over the network via a protocol like HTTP.
+    encoded_swallow: str = base64.urlsafe_b64encode(b"unladen swallow").decode("utf-8")
+    assert type(encoded_swallow) == str
+    foo.baz = encoded_swallow
+    assert foo.baz == b"unladen swallow"
+
+
+def test_bytes_to_dict_bidi():
+    class Foo(proto.Message):
+        bar = proto.Field(proto.BYTES, number=1)
+
+    foo = Foo(bar=b"spam")
+
+    foo_dict = Foo.to_dict(foo)
+    foo_two = Foo(foo_dict)
+
+    assert foo == foo_two

tests/test_marshal_register.py

@@ -33,19 +33,6 @@ def to_python(self, value, *, absent=None):
     assert isinstance(marshal._rules[empty_pb2.Empty], Rule)
 
 
-def test_invalid_target_registration():
-    marshal = BaseMarshal()
-    with pytest.raises(TypeError):
-
-        @marshal.register(object)
-        class Rule:
-            def to_proto(self, value):
-                return value
-
-            def to_python(self, value, *, absent=None):
-                return value
-
-
 def test_invalid_marshal_class():
     marshal = BaseMarshal()
     with pytest.raises(TypeError):