Support Proto2 extensions in JSON coder/decoder

Protos/unittest_swift_extension.proto

@@ -53,6 +53,7 @@ message Msg2 {
 
 extend Msg1 {
   optional int32 a_b = 1;
+  optional Msg2 m2 = 2;
 }
 
 extend Msg2 {

Reference/unittest_swift_extension.pb.swift

@@ -236,6 +236,21 @@ extension ProtobufUnittest_Extend_Msg1 {
   mutating func clearProtobufUnittest_Extend_aB() {
     clearExtensionValue(ext: ProtobufUnittest_Extend_Extensions_a_b)
   }
+
+  var ProtobufUnittest_Extend_m2: ProtobufUnittest_Extend_Msg2 {
+    get {return getExtensionValue(ext: ProtobufUnittest_Extend_Extensions_m2) ?? ProtobufUnittest_Extend_Msg2()}
+    set {setExtensionValue(ext: ProtobufUnittest_Extend_Extensions_m2, value: newValue)}
+  }
+  /// Returns true if extension `ProtobufUnittest_Extend_Extensions_m2`
+  /// has been explicitly set.
+  var hasProtobufUnittest_Extend_m2: Bool {
+    return hasExtensionValue(ext: ProtobufUnittest_Extend_Extensions_m2)
+  }
+  /// Clears the value of extension `ProtobufUnittest_Extend_Extensions_m2`.
+  /// Subsequent reads from it will return its default value.
+  mutating func clearProtobufUnittest_Extend_m2() {
+    clearExtensionValue(ext: ProtobufUnittest_Extend_Extensions_m2)
+  }
 }
 
 extension ProtobufUnittest_Extend_Msg2 {
@@ -333,6 +348,7 @@ let ProtobufUnittest_Extend_UnittestSwiftExtension_Extensions: SwiftProtobuf.Sim
   ProtobufUnittest_Extend_Extensions_b,
   ProtobufUnittest_Extend_Extensions_C,
   ProtobufUnittest_Extend_Extensions_a_b,
+  ProtobufUnittest_Extend_Extensions_m2,
   ProtobufUnittest_Extend_Extensions_aB,
   ProtobufUnittest_Extend_Extensions_ext_a,
   ProtobufUnittest_Extend_Extensions_ext_b,
@@ -359,6 +375,11 @@ let ProtobufUnittest_Extend_Extensions_a_b = SwiftProtobuf.MessageExtension<Swif
   fieldName: "protobuf_unittest.extend.a_b"
 )
 
+let ProtobufUnittest_Extend_Extensions_m2 = SwiftProtobuf.MessageExtension<SwiftProtobuf.OptionalMessageExtensionField<ProtobufUnittest_Extend_Msg2>, ProtobufUnittest_Extend_Msg1>(
+  _protobuf_fieldNumber: 2,
+  fieldName: "protobuf_unittest.extend.m2"
+)
+
 let ProtobufUnittest_Extend_Extensions_aB = SwiftProtobuf.MessageExtension<SwiftProtobuf.OptionalExtensionField<SwiftProtobuf.ProtobufInt32>, ProtobufUnittest_Extend_Msg2>(
   _protobuf_fieldNumber: 1,
   fieldName: "protobuf_unittest.extend.aB"

Sources/SwiftProtobuf/JSONDecoder.swift

@@ -27,13 +27,6 @@ internal struct JSONDecoder: Decoder {
     throw JSONDecodingError.conflictingOneOf
   }
 
-  internal init(source: UnsafeRawBufferPointer, options: JSONDecodingOptions) {
-    self.options = options
-    self.scanner = JSONScanner(source: source,
-                               messageDepthLimit: self.options.messageDepthLimit,
-                               ignoreUnknownFields: self.options.ignoreUnknownFields)
-  }
-
   internal init(source: UnsafeRawBufferPointer, options: JSONDecodingOptions,
                 messageType: Message.Type, extensions: ExtensionMap?) {
     self.options = options
@@ -44,11 +37,12 @@ internal struct JSONDecoder: Decoder {
     self.extensions = extensions
   }
 
-  private init(decoder: JSONDecoder) {
+  private init(decoder: JSONDecoder, messageType: Message.Type) {
     // The scanner is copied over along with the options.
-    scanner = decoder.scanner
-    options = decoder.options
-    extensions = decoder.extensions
+    self.scanner = decoder.scanner
+    self.options = decoder.options
+    self.extensions = decoder.extensions
+    self.messageType = messageType
   }
 
   mutating func nextFieldNumber() throws -> Int? {
@@ -545,7 +539,7 @@ internal struct JSONDecoder: Decoder {
     if value == nil {
       value = M()
     }
-    var subDecoder = JSONDecoder(decoder: self)
+    var subDecoder = JSONDecoder(decoder: self, messageType: M.self)
     try subDecoder.decodeFullObject(message: &value!)
     assert(scanner.recursionBudget == subDecoder.scanner.recursionBudget)
     scanner = subDecoder.scanner
@@ -576,7 +570,7 @@ internal struct JSONDecoder: Decoder {
         }
       } else {
         var message = M()
-        var subDecoder = JSONDecoder(decoder: self)
+        var subDecoder = JSONDecoder(decoder: self, messageType: M.self)
         try subDecoder.decodeFullObject(message: &message)
         value.append(message)
         assert(scanner.recursionBudget == subDecoder.scanner.recursionBudget)
@@ -713,19 +707,20 @@ internal struct JSONDecoder: Decoder {
     messageType: Message.Type,
     fieldNumber: Int
   ) throws {
-    if let ext = extensions?[messageType, fieldNumber] {
-      var fieldValue = values[fieldNumber]
-      if fieldValue != nil {
-        try fieldValue!.decodeExtensionField(decoder: &self)
-      } else {
-        fieldValue = try ext._protobuf_newField(decoder: &self)
-      }
-      if fieldValue != nil {
-        values[fieldNumber] = fieldValue
-      } else {
-        // This most likely indicates a bug in our extension support.
-        throw TextFormatDecodingError.internalExtensionError
-      }
+    guard let ext = extensions?[messageType, fieldNumber] else {
+      return
+    }
+    var fieldValue = values[fieldNumber]
+    if fieldValue != nil {
+      try fieldValue!.decodeExtensionField(decoder: &self)
+    } else {
+      fieldValue = try ext._protobuf_newField(decoder: &self)
+    }
+    if fieldValue != nil {
+      values[fieldNumber] = fieldValue
+    } else {
+      // This most likely indicates a bug in our extension support.
+      throw TextFormatDecodingError.internalExtensionError
     }
   }
 }

Sources/SwiftProtobuf/JSONEncodingVisitor.swift

@@ -173,10 +173,17 @@ internal struct JSONEncodingVisitor: Visitor {
       encoder.append(text: json)
     } else if let newNameMap = (M.self as? _ProtoNameProviding.Type)?._protobuf_nameMap {
       encoder.startNestedObject()
+      // Preserve outer object's name and extension maps; restore them before returning
       let oldNameMap = self.nameMap
+      let oldExtensions = self.extensions
+      defer {
+        self.nameMap = oldNameMap
+        self.extensions = oldExtensions
+      }
+      // Install inner object's name and extension maps
       self.nameMap = newNameMap
+      self.extensions = (value as? ExtensibleMessage)?._protobuf_extensionFieldValues
       try value.traverse(visitor: &self)
-      self.nameMap = oldNameMap
       endObject()
     } else {
       throw JSONEncodingError.missingFieldNames
@@ -300,17 +307,22 @@ internal struct JSONEncodingVisitor: Visitor {
         encoder.append(text: json)
       }
     } else if let newNameMap = (M.self as? _ProtoNameProviding.Type)?._protobuf_nameMap {
-      // Install inner object's name map
+      // Preserve name and extension maps for outer object
       let oldNameMap = self.nameMap
+      let oldExtensions = self.extensions
+      // Restore outer object's name and extension maps before returning
+      defer {
+        self.nameMap = oldNameMap
+        self.extensions = oldExtensions
+      }
       self.nameMap = newNameMap
-      // Restore outer object's name map before returning
-      defer { self.nameMap = oldNameMap }
       for v in value {
         if comma {
           encoder.comma()
         }
         comma = true
         encoder.startNestedObject()
+        self.extensions = (v as? ExtensibleMessage)?._protobuf_extensionFieldValues
         try v.traverse(visitor: &self)
         encoder.endObject()
       }

Sources/SwiftProtobuf/Message+JSONAdditions.swift

@@ -70,14 +70,7 @@ extension Message {
     jsonString: String,
     options: JSONDecodingOptions = JSONDecodingOptions()
   ) throws {
-    if jsonString.isEmpty {
-      throw JSONDecodingError.truncated
-    }
-    if let data = jsonString.data(using: String.Encoding.utf8) {
-      try self.init(jsonUTF8Data: data, options: options)
-    } else {
-      throw JSONDecodingError.truncated
-    }
+    try self.init(jsonString: jsonString, extensions: nil, options: options)
   }
 
   /// Creates a new message by decoding the given string containing a
@@ -114,27 +107,7 @@ extension Message {
     jsonUTF8Data: Data,
     options: JSONDecodingOptions = JSONDecodingOptions()
   ) throws {
-    self.init()
-    try jsonUTF8Data.withUnsafeBytes { (body: UnsafeRawBufferPointer) in
-      // Empty input is valid for binary, but not for JSON.
-      guard body.count > 0 else {
-        throw JSONDecodingError.truncated
-      }
-      var decoder = JSONDecoder(source: body, options: options)
-      if decoder.scanner.skipOptionalNull() {
-        if let customCodable = Self.self as? _CustomJSONCodable.Type,
-           let message = try customCodable.decodedFromJSONNull() {
-          self = message as! Self
-        } else {
-          throw JSONDecodingError.illegalNull
-        }
-      } else {
-        try decoder.decodeFullObject(message: &self)
-      }
-      if !decoder.scanner.complete {
-        throw JSONDecodingError.trailingGarbage
-      }
-    }
+    try self.init(jsonUTF8Data: jsonUTF8Data, extensions: nil, options: options)
   }
 
   /// Creates a new message by decoding the given `Data` containing a

Sources/SwiftProtobuf/Message+JSONArrayAdditions.swift

@@ -95,7 +95,8 @@ extension Message {
       var array = [Self]()
 
       if body.count > 0 {
-        var decoder = JSONDecoder(source: body, options: options)
+        var decoder = JSONDecoder(source: body, options: options,
+                                  messageType: Self.self, extensions: nil)
         try decoder.decodeRepeatedMessageField(value: &array)
         if !decoder.scanner.complete {
           throw JSONDecodingError.trailingGarbage

Tests/LinuxMain.swift

@@ -649,6 +649,12 @@ extension Test_JSON_Extensions {
     ]
 }
 
+extension Test_JSON_RecursiveNested_Extensions {
+    static var allTests = [
+        ("test_nestedMessage", test_nestedMessage)
+    ]
+}
+
 extension Test_JSON_Group {
     static var allTests = [
         ("testOptionalGroup", testOptionalGroup),
@@ -1193,6 +1199,7 @@ XCTMain(
         testCase(Test_JSON_Array.allTests),
         testCase(Test_JSON_Conformance.allTests),
         testCase(Test_JSON_Extensions.allTests),
+        testCase(Test_JSON_RecursiveNested_Extensions.allTests),
         testCase(Test_JSON_Group.allTests),
         testCase(Test_Map.allTests),
         testCase(Test_MapFields_Access_Proto2.allTests),

Tests/SwiftProtobufTests/Test_JSON_Extensions.swift

@@ -41,5 +41,22 @@ class Test_JSON_Extensions: XCTestCase, PBTestHelpers {
             o.ProtobufUnittest_optionalInt32Extension = 17
         }
     }
+}
 
+class Test_JSON_RecursiveNested_Extensions: XCTestCase, PBTestHelpers {
+    typealias MessageTestType = ProtobufUnittest_Extend_Msg1
+    let extensions = ProtobufUnittest_Extend_UnittestSwiftExtension_Extensions
+
+    func test_nestedMessage() throws {
+        assertJSONEncode("{\"[protobuf_unittest.extend.a_b]\":12}",
+                         extensions: extensions) {
+                            (o: inout MessageTestType) in
+                            o.ProtobufUnittest_Extend_aB = 12
+        }
+
+        assertJSONDecodeSucceeds("{\"[protobuf_unittest.extend.m2]\":{\"[protobuf_unittest.extend.aB]\":23}}", extensions: extensions) {
+            $0.ProtobufUnittest_Extend_m2.ProtobufUnittest_Extend_aB == 23
+        }
+    }
+
 }

Tests/SwiftProtobufTests/unittest_swift_extension.pb.swift

@@ -236,6 +236,21 @@ extension ProtobufUnittest_Extend_Msg1 {
   mutating func clearProtobufUnittest_Extend_aB() {
     clearExtensionValue(ext: ProtobufUnittest_Extend_Extensions_a_b)
   }
+
+  var ProtobufUnittest_Extend_m2: ProtobufUnittest_Extend_Msg2 {
+    get {return getExtensionValue(ext: ProtobufUnittest_Extend_Extensions_m2) ?? ProtobufUnittest_Extend_Msg2()}
+    set {setExtensionValue(ext: ProtobufUnittest_Extend_Extensions_m2, value: newValue)}
+  }
+  /// Returns true if extension `ProtobufUnittest_Extend_Extensions_m2`
+  /// has been explicitly set.
+  var hasProtobufUnittest_Extend_m2: Bool {
+    return hasExtensionValue(ext: ProtobufUnittest_Extend_Extensions_m2)
+  }
+  /// Clears the value of extension `ProtobufUnittest_Extend_Extensions_m2`.
+  /// Subsequent reads from it will return its default value.
+  mutating func clearProtobufUnittest_Extend_m2() {
+    clearExtensionValue(ext: ProtobufUnittest_Extend_Extensions_m2)
+  }
 }
 
 extension ProtobufUnittest_Extend_Msg2 {
@@ -333,6 +348,7 @@ let ProtobufUnittest_Extend_UnittestSwiftExtension_Extensions: SwiftProtobuf.Sim
   ProtobufUnittest_Extend_Extensions_b,
   ProtobufUnittest_Extend_Extensions_C,
   ProtobufUnittest_Extend_Extensions_a_b,
+  ProtobufUnittest_Extend_Extensions_m2,
   ProtobufUnittest_Extend_Extensions_aB,
   ProtobufUnittest_Extend_Extensions_ext_a,
   ProtobufUnittest_Extend_Extensions_ext_b,
@@ -359,6 +375,11 @@ let ProtobufUnittest_Extend_Extensions_a_b = SwiftProtobuf.MessageExtension<Swif
   fieldName: "protobuf_unittest.extend.a_b"
 )
 
+let ProtobufUnittest_Extend_Extensions_m2 = SwiftProtobuf.MessageExtension<SwiftProtobuf.OptionalMessageExtensionField<ProtobufUnittest_Extend_Msg2>, ProtobufUnittest_Extend_Msg1>(
+  _protobuf_fieldNumber: 2,
+  fieldName: "protobuf_unittest.extend.m2"
+)
+
 let ProtobufUnittest_Extend_Extensions_aB = SwiftProtobuf.MessageExtension<SwiftProtobuf.OptionalExtensionField<SwiftProtobuf.ProtobufInt32>, ProtobufUnittest_Extend_Msg2>(
   _protobuf_fieldNumber: 1,
   fieldName: "protobuf_unittest.extend.aB"