From 9e5d5bd5f9783d00ae5da31f0462f741a5cf0307 Mon Sep 17 00:00:00 2001
From: Andrew Kozlik <andrew.kozlik@gmail.com>
Date: Wed, 7 Oct 2020 18:18:03 +0200
Subject: [PATCH] feat: Add support for experimental field flag in protobuf.

---
 common/protob/messages.proto     | 6 ++++++
 common/protob/pb2py              | 5 +++++
 core/src/protobuf.py             | 6 ++++++
 python/src/trezorlib/protobuf.py | 3 +++
 4 files changed, 20 insertions(+)

diff --git a/common/protob/messages.proto b/common/protob/messages.proto
index a4f922b50..d6675963a 100644
--- a/common/protob/messages.proto
+++ b/common/protob/messages.proto
@@ -32,6 +32,12 @@ extend google.protobuf.MessageOptions {
 }
 
 
+/** Options for tagging field types */
+extend google.protobuf.FieldOptions {
+    optional bool experimental = 52001;   // indicate that a field is intended for development and beta testing only
+}
+
+
 /**
  * Mapping between Trezor wire identifier (uint) and a protobuf message
  */
diff --git a/common/protob/pb2py b/common/protob/pb2py
index 8b8196df4..2b457ec98 100755
--- a/common/protob/pb2py
+++ b/common/protob/pb2py
@@ -65,6 +65,7 @@ class ProtoField:
     orig = attr.ib()
     repeated = attr.ib()
     required = attr.ib()
+    experimental = attr.ib()
     type_name = attr.ib()
     proto_type = attr.ib()
     py_type = attr.ib()
@@ -78,6 +79,7 @@ class ProtoField:
     def from_field(cls, descriptor, field):
         repeated = field.label == field.LABEL_REPEATED
         required = field.label == field.LABEL_REQUIRED
+        experimental = bool(descriptor._get_extension(field, "experimental"))
         # ignore package path
         type_name = field.type_name.rsplit(".")[-1]
 
@@ -115,6 +117,7 @@ class ProtoField:
             orig=field,
             repeated=repeated,
             required=required,
+            experimental=experimental,
             type_name=type_name,
             proto_type=proto_type,
             py_type=py_type,
@@ -306,6 +309,8 @@ class Descriptor:
                 flags = "p.FLAG_REPEATED"
             elif field.required:
                 flags = "p.FLAG_REQUIRED"
+            elif field.experimental:
+                flags = "p.FLAG_EXPERIMENTAL"
             else:
                 flags = field.default_value
 
diff --git a/core/src/protobuf.py b/core/src/protobuf.py
index 6e42aa785..9ae3d97b1 100644
--- a/core/src/protobuf.py
+++ b/core/src/protobuf.py
@@ -179,6 +179,7 @@ class LimitedReader:
 
 FLAG_REPEATED = object()
 FLAG_REQUIRED = object()
+FLAG_EXPERIMENTAL = object()
 
 if False:
     MessageTypeDef = Union[
@@ -220,6 +221,8 @@ def load_message(
     for fname, _, fdefault in fields.values():
         if fdefault is FLAG_REPEATED:
             fdefault = []
+        elif fdefault is FLAG_EXPERIMENTAL:
+            fdefault = None
         setattr(msg, fname, fdefault)
 
     if False:
@@ -252,6 +255,9 @@ def load_message(
         if wtype != ftype.WIRE_TYPE:
             raise TypeError  # parsed wire type differs from the schema
 
+        if fdefault is FLAG_EXPERIMENTAL and not experimental_enabled:
+            raise ValueError  # experimental fields not enabled
+
         ivalue = load_uvarint(reader)
 
         if ftype is UVarintType:
diff --git a/python/src/trezorlib/protobuf.py b/python/src/trezorlib/protobuf.py
index c38806331..22b050af7 100644
--- a/python/src/trezorlib/protobuf.py
+++ b/python/src/trezorlib/protobuf.py
@@ -280,6 +280,7 @@ class CountingWriter:
 
 FLAG_REPEATED = object()
 FLAG_REQUIRED = object()
+FLAG_EXPERIMENTAL = object()
 
 
 def decode_packed_array_field(ftype: FieldType, reader: Reader) -> List[Any]:
@@ -334,6 +335,8 @@ def load_message(reader: Reader, msg_type: Type[MT]) -> MT:
     for fname, _, fdefault in fields.values():
         if fdefault is FLAG_REPEATED:
             msg_dict[fname] = []
+        elif fdefault is FLAG_EXPERIMENTAL:
+            msg_dict[fname] = None
         elif fdefault is not FLAG_REQUIRED:
             msg_dict[fname] = fdefault