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/Conformance/failure_list_swift.txt

@@ -1,2 +1 @@
-# Tracking in https://github.com/apple/swift-protobuf/issues/993
-Recommended.Proto2.JsonInput.FieldNameExtension.Validator
+# No failures

Sources/Conformance/main.swift

@@ -85,15 +85,18 @@ func buildResponse(serializedData: Data) -> Conformance_ConformanceResponse {
     }
 
     let msgType: SwiftProtobuf.Message.Type
+    let extensions: SwiftProtobuf.ExtensionMap
     switch request.messageType {
     case "":
         // Note: This case is here to cover using a old version of the conformance test
         // runner that don't know about this field, and it is thus implicit.
         fallthrough
     case ProtobufTestMessages_Proto3_TestAllTypesProto3.protoMessageName:
         msgType = ProtobufTestMessages_Proto3_TestAllTypesProto3.self
+        extensions = SwiftProtobuf.SimpleExtensionMap()
     case ProtobufTestMessages_Proto2_TestAllTypesProto2.protoMessageName:
         msgType = ProtobufTestMessages_Proto2_TestAllTypesProto2.self
+        extensions = ProtobufTestMessages_Proto2_TestMessagesProto2_Extensions
     default:
         response.runtimeError = "Unexpected message type: \(request.messageType)"
         return response
@@ -103,7 +106,7 @@ func buildResponse(serializedData: Data) -> Conformance_ConformanceResponse {
     switch request.payload {
     case .protobufPayload(let data)?:
         do {
-            testMessage = try msgType.init(serializedData: data)
+            testMessage = try msgType.init(serializedData: data, extensions: extensions)
         } catch let e {
             response.parseError = "Protobuf failed to parse: \(e)"
             return response
@@ -112,7 +115,9 @@ func buildResponse(serializedData: Data) -> Conformance_ConformanceResponse {
         var options = JSONDecodingOptions()
         options.ignoreUnknownFields = (request.testCategory == .jsonIgnoreUnknownParsingTest)
         do {
-            testMessage = try msgType.init(jsonString: json, options: options)
+            testMessage = try msgType.init(jsonString: json,
+                                           extensions: extensions,
+                                           options: options)
         } catch let e {
             response.parseError = "JSON failed to parse: \(e)"
             return response
@@ -124,7 +129,7 @@ func buildResponse(serializedData: Data) -> Conformance_ConformanceResponse {
         return response
     case .textPayload(let textFormat)?:
         do {
-            testMessage = try msgType.init(textFormatString: textFormat)
+            testMessage = try msgType.init(textFormatString: textFormat, extensions: extensions)
         } catch let e {
             response.parseError = "Protobuf failed to parse: \(e)"
             return response

Sources/SwiftProtobuf/AnyMessageStorage.swift

@@ -25,8 +25,8 @@ fileprivate func serializeAnyJSON(
   typeURL: String,
   options: JSONEncodingOptions
 ) throws -> String {
-  var visitor = try JSONEncodingVisitor(message: message, options: options)
-  visitor.startObject()
+  var visitor = try JSONEncodingVisitor(type: type(of: message), options: options)
+  visitor.startObject(message: message)
   visitor.encodeField(name: "@type", stringValue: typeURL)
   if let m = message as? _CustomJSONCodable {
     let value = try m.encodedJSONString(options: options)
@@ -58,19 +58,18 @@ fileprivate func asJSONObject(body: Data) -> Data {
 }
 
 fileprivate func unpack(contentJSON: Data,
+                        extensions: ExtensionMap,
                         options: JSONDecodingOptions,
                         as messageType: Message.Type) throws -> Message {
   guard messageType is _CustomJSONCodable.Type else {
     let contentJSONAsObject = asJSONObject(body: contentJSON)
-    return try messageType.init(jsonUTF8Data: contentJSONAsObject, options: options)
+    return try messageType.init(jsonUTF8Data: contentJSONAsObject, extensions: extensions, options: options)
   }
 
   var value = String()
   try contentJSON.withUnsafeBytes { (body: UnsafeRawBufferPointer) in
     if body.count > 0 {
-      var scanner = JSONScanner(source: body,
-                                messageDepthLimit: options.messageDepthLimit,
-                                ignoreUnknownFields: options.ignoreUnknownFields)
+      var scanner = JSONScanner(source: body, options: options, extensions: extensions)
       let key = try scanner.nextQuotedString()
       if key != "value" {
         // The only thing within a WKT should be "value".
@@ -85,7 +84,7 @@ fileprivate func unpack(contentJSON: Data,
       }
     }
   }
-  return try messageType.init(jsonString: value, options: options)
+  return try messageType.init(jsonString: value, extensions: extensions, options: options)
 }
 
 internal class AnyMessageStorage {
@@ -109,6 +108,7 @@ internal class AnyMessageStorage {
         }
         do {
           let m = try unpack(contentJSON: contentJSON,
+                             extensions: SimpleExtensionMap(),
                              options: options,
                              as: messageType)
           return try m.serializedData(partial: true)
@@ -154,7 +154,7 @@ internal class AnyMessageStorage {
     return encodedType == M.protoMessageName
   }
 
-  // This is only ever called with the expactation that target will be fully
+  // This is only ever called with the expectation that target will be fully
   // replaced during the unpacking and never as a merge.
   func unpackTo<M: Message>(
     target: inout M,
@@ -181,6 +181,7 @@ internal class AnyMessageStorage {
 
     case .contentJSON(let contentJSON, let options):
       target = try unpack(contentJSON: contentJSON,
+                          extensions: extensions ?? SimpleExtensionMap(),
                           options: options,
                           as: M.self) as! M
     }
@@ -202,11 +203,21 @@ internal class AnyMessageStorage {
       // never inserted.
       break
 
-    case .contentJSON:
-      // contentJSON requires a good URL and our ability to look up
-      // the message type to transcode.
-      if Google_Protobuf_Any.messageType(forTypeURL: _typeURL) == nil {
-        // Isn't registered, we can't transform it for binary.
+    case .contentJSON(let contentJSON, let options):
+      // contentJSON requires we have the type available for decoding
+      guard let messageType = Google_Protobuf_Any.messageType(forTypeURL: _typeURL) else {
+          throw BinaryEncodingError.anyTranscodeFailure
+      }
+      do {
+        // Decodes the full JSON and then discard the result.
+        // The regular traversal will decode this again by querying the
+        // `value` field, but that has no way to fail.  As a result,
+        // we need this to accurately handle decode errors.
+        _ = try unpack(contentJSON: contentJSON,
+                       extensions: SimpleExtensionMap(),
+                       options: options,
+                       as: messageType)
+      } catch {
         throw BinaryEncodingError.anyTranscodeFailure
       }
     }
@@ -271,6 +282,7 @@ extension AnyMessageStorage {
       if let messageType = Google_Protobuf_Any.messageType(forTypeURL: _typeURL) {
         do {
           let m = try unpack(contentJSON: contentJSON,
+                             extensions: SimpleExtensionMap(),
                              options: options,
                              as: messageType)
           emitVerboseTextForm(visitor: &visitor, message: m, typeURL: _typeURL)

Sources/SwiftProtobuf/JSONDecoder.swift

@@ -16,26 +16,30 @@ import Foundation
 
 internal struct JSONDecoder: Decoder {
   internal var scanner: JSONScanner
-  internal var options: JSONDecodingOptions
+  internal var messageType: Message.Type
   private var fieldCount = 0
   private var isMapKey = false
   private var fieldNameMap: _NameMap?
 
+  internal var options: JSONDecodingOptions {
+    return scanner.options
+  }
+
   mutating func handleConflictingOneOf() throws {
     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?) {
+    let scanner = JSONScanner(source: source,
+                               options: options,
+                               extensions: extensions)
+    self.init(scanner: scanner, messageType: messageType)
   }
 
-  private init(decoder: JSONDecoder) {
-    // The scanner is copied over along with the options.
-    scanner = decoder.scanner
-    options = decoder.options
+  private init(scanner: JSONScanner, messageType: Message.Type) {
+    self.scanner = scanner
+    self.messageType = messageType
   }
 
   mutating func nextFieldNumber() throws -> Int? {
@@ -45,7 +49,9 @@ internal struct JSONDecoder: Decoder {
     if fieldCount > 0 {
       try scanner.skipRequiredComma()
     }
-    if let fieldNumber = try scanner.nextFieldNumber(names: fieldNameMap!) {
+    let fieldNumber = try scanner.nextFieldNumber(names: fieldNameMap!,
+                                                  messageType: messageType)
+    if let fieldNumber = fieldNumber {
       fieldCount += 1
       return fieldNumber
     }
@@ -529,7 +535,7 @@ internal struct JSONDecoder: Decoder {
     if value == nil {
       value = M()
     }
-    var subDecoder = JSONDecoder(decoder: self)
+    var subDecoder = JSONDecoder(scanner: scanner, messageType: M.self)
     try subDecoder.decodeFullObject(message: &value!)
     assert(scanner.recursionBudget == subDecoder.scanner.recursionBudget)
     scanner = subDecoder.scanner
@@ -560,7 +566,7 @@ internal struct JSONDecoder: Decoder {
         }
       } else {
         var message = M()
-        var subDecoder = JSONDecoder(decoder: self)
+        var subDecoder = JSONDecoder(scanner: scanner, messageType: M.self)
         try subDecoder.decodeFullObject(message: &message)
         value.append(message)
         assert(scanner.recursionBudget == subDecoder.scanner.recursionBudget)
@@ -697,6 +703,15 @@ internal struct JSONDecoder: Decoder {
     messageType: Message.Type,
     fieldNumber: Int
   ) throws {
-    throw JSONDecodingError.schemaMismatch
+    // Force-unwrap: we can only get here if the extension exists.
+    let ext = scanner.extensions[messageType, fieldNumber]!
+
+    var fieldValue = values[fieldNumber]
+    if fieldValue != nil {
+      try fieldValue!.decodeExtensionField(decoder: &self)
+    } else {
+      fieldValue = try ext._protobuf_newField(decoder: &self)
+    }
+    values[fieldNumber] = fieldValue!
   }
 }

Sources/SwiftProtobuf/JSONEncoder.swift

@@ -128,6 +128,17 @@ internal struct JSONEncoder {
         separator = asciiComma
     }
 
+    /// Begin a new extension field
+    internal mutating func startExtensionField(name: String) {
+        if let s = separator {
+            data.append(s)
+        }
+        append(staticText: "\"[")
+        data.append(contentsOf: name.utf8)
+        append(staticText: "]\":")
+        separator = asciiComma
+    }
+
     /// Append an open square bracket `[` to the JSON.
     internal mutating func startArray() {
         data.append(asciiOpenSquareBracket)
@@ -146,15 +157,17 @@ internal struct JSONEncoder {
     }
 
     /// Append an open curly brace `{` to the JSON.
-    internal mutating func startObject() {
+    /// Assumes this object is part of an array of objects.
+    internal mutating func startArrayObject() {
         if let s = separator {
             data.append(s)
         }
         data.append(asciiOpenCurlyBracket)
         separator = nil
     }
 
-    internal mutating func startNestedObject() {
+    /// Append an open curly brace `{` to the JSON.
+    internal mutating func startObject() {
         data.append(asciiOpenCurlyBracket)
         separator = nil
     }