// Protocol Buffers - Google's data interchange format
// Copyright 2008 Google Inc.  All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd

#import "GPBCodedInputStream.h"
#import "GPBMessage_PackagePrivate.h"
#import "GPBTestUtilities.h"
#import "GPBUnknownField.h"
#import "GPBUnknownField_PackagePrivate.h"
#import "GPBUnknownFields.h"
#import "GPBWireFormat.h"
#import "objectivec/Tests/Unittest.pbobjc.h"
#import "objectivec/Tests/UnittestMset.pbobjc.h"

@interface WireFormatTests : GPBTestCase
@end

@implementation WireFormatTests

- (void)testSerialization {
  TestAllTypes* message = [self allSetRepeatedCount:kGPBDefaultRepeatCount];

  NSData* rawBytes = message.data;
  [self assertFieldsInOrder:rawBytes];
  XCTAssertEqual(message.serializedSize, (size_t)rawBytes.length);

  TestAllTypes* message2 = [TestAllTypes parseFromData:rawBytes error:NULL];

  [self assertAllFieldsSet:message2 repeatedCount:kGPBDefaultRepeatCount];
}

- (void)testSerializationPacked {
  TestPackedTypes* message = [self packedSetRepeatedCount:kGPBDefaultRepeatCount];

  NSData* rawBytes = message.data;
  [self assertFieldsInOrder:rawBytes];
  XCTAssertEqual(message.serializedSize, (size_t)rawBytes.length);

  TestPackedTypes* message2 = [TestPackedTypes parseFromData:rawBytes error:NULL];

  [self assertPackedFieldsSet:message2 repeatedCount:kGPBDefaultRepeatCount];
}

- (void)testSerializeExtensions {
  // TestAllTypes and TestAllExtensions should have compatible wire formats,
  // so if we serealize a TestAllExtensions then parse it as TestAllTypes
  // it should work.

  TestAllExtensions* message = [self allExtensionsSetRepeatedCount:kGPBDefaultRepeatCount];
  NSData* rawBytes = message.data;
  [self assertFieldsInOrder:rawBytes];
  XCTAssertEqual(message.serializedSize, (size_t)rawBytes.length);

  TestAllTypes* message2 = [TestAllTypes parseFromData:rawBytes error:NULL];

  [self assertAllFieldsSet:message2 repeatedCount:kGPBDefaultRepeatCount];
}

- (void)testSerializePackedExtensions {
  // TestPackedTypes and TestPackedExtensions should have compatible wire
  // formats; check that they serialize to the same string.
  TestPackedExtensions* message = [self packedExtensionsSetRepeatedCount:kGPBDefaultRepeatCount];
  NSData* rawBytes = message.data;
  [self assertFieldsInOrder:rawBytes];

  TestPackedTypes* message2 = [self packedSetRepeatedCount:kGPBDefaultRepeatCount];
  NSData* rawBytes2 = message2.data;

  XCTAssertEqualObjects(rawBytes, rawBytes2);
}

- (void)testParseExtensions {
  // TestAllTypes and TestAllExtensions should have compatible wire formats,
  // so if we serialize a TestAllTypes then parse it as TestAllExtensions
  // it should work.

  TestAllTypes* message = [self allSetRepeatedCount:kGPBDefaultRepeatCount];
  NSData* rawBytes = message.data;
  [self assertFieldsInOrder:rawBytes];

  GPBExtensionRegistry* registry = [self extensionRegistry];

  TestAllExtensions* message2 = [TestAllExtensions parseFromData:rawBytes
                                               extensionRegistry:registry
                                                           error:NULL];

  [self assertAllExtensionsSet:message2 repeatedCount:kGPBDefaultRepeatCount];
}

- (void)testExtensionsSerializedSize {
  size_t allSet = [self allSetRepeatedCount:kGPBDefaultRepeatCount].serializedSize;
  size_t extensionSet = [self allExtensionsSetRepeatedCount:kGPBDefaultRepeatCount].serializedSize;
  XCTAssertEqual(allSet, extensionSet);
}

- (void)testParsePackedExtensions {
  // Ensure that packed extensions can be properly parsed.
  TestPackedExtensions* message = [self packedExtensionsSetRepeatedCount:kGPBDefaultRepeatCount];
  NSData* rawBytes = message.data;
  [self assertFieldsInOrder:rawBytes];

  GPBExtensionRegistry* registry = [self extensionRegistry];

  TestPackedExtensions* message2 = [TestPackedExtensions parseFromData:rawBytes
                                                     extensionRegistry:registry
                                                                 error:NULL];

  [self assertPackedExtensionsSet:message2 repeatedCount:kGPBDefaultRepeatCount];
}

const int kUnknownTypeId = 1550055;
const int kUnknownTypeId2 = 1550056;

- (void)testSerializeMessageSet {
  // Set up a MSetMessage with two known messages and an unknown one.
  MSetMessage* message_set = [MSetMessage message];
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  [[message_set getExtension:MSetMessageExtension1_extension_MessageSetExtension()] setI:123];
  [[message_set getExtension:MSetMessageExtension2_extension_MessageSetExtension()] setStr:@"foo"];
#else
  [[message_set getExtension:[MSetMessageExtension1 messageSetExtension]] setI:123];
  [[message_set getExtension:[MSetMessageExtension2 messageSetExtension]] setStr:@"foo"];
#endif

  GPBUnknownFields* ufs = [[[GPBUnknownFields alloc] init] autorelease];
  GPBUnknownFields* group = [ufs addGroupWithFieldNumber:GPBWireFormatMessageSetItem];
  [group addFieldNumber:GPBWireFormatMessageSetTypeId varint:kUnknownTypeId2];
  [group addFieldNumber:GPBWireFormatMessageSetMessage lengthDelimited:DataFromCStr("baz")];
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  XCTAssertTrue([message_set
      mergeUnknownFields:ufs
       extensionRegistry:MSet_Objc_Protobuf_Tests_Mset_MSetUnittestMsetRoot_Registry()
                   error:NULL]);
#else
  XCTAssertTrue([message_set mergeUnknownFields:ufs
                              extensionRegistry:[MSetUnittestMsetRoot extensionRegistry]
                                          error:NULL]);
#endif

  NSData* data = [message_set data];

  // Parse back using MSetRawMessageSet and check the contents.
  MSetRawMessageSet* raw = [MSetRawMessageSet parseFromData:data error:NULL];

  GPBUnknownFields* ufs2 = [[[GPBUnknownFields alloc] initFromMessage:raw] autorelease];
  XCTAssertTrue(ufs2.empty);

  XCTAssertEqual(raw.itemArray.count, (NSUInteger)3);
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  XCTAssertEqual((uint32_t)[raw.itemArray[0] typeId],
                 MSetMessageExtension1_extension_MessageSetExtension().fieldNumber);
  XCTAssertEqual((uint32_t)[raw.itemArray[1] typeId],
                 MSetMessageExtension2_extension_MessageSetExtension().fieldNumber);
#else
  XCTAssertEqual((uint32_t)[raw.itemArray[0] typeId],
                 [MSetMessageExtension1 messageSetExtension].fieldNumber);
  XCTAssertEqual((uint32_t)[raw.itemArray[1] typeId],
                 [MSetMessageExtension2 messageSetExtension].fieldNumber);
#endif
  XCTAssertEqual([raw.itemArray[2] typeId], kUnknownTypeId2);

  MSetMessageExtension1* message1 =
      [MSetMessageExtension1 parseFromData:[((MSetRawMessageSet_Item*)raw.itemArray[0]) message]
                                     error:NULL];
  XCTAssertEqual(message1.i, 123);

  MSetMessageExtension2* message2 =
      [MSetMessageExtension2 parseFromData:[((MSetRawMessageSet_Item*)raw.itemArray[1]) message]
                                     error:NULL];
  XCTAssertEqualObjects(message2.str, @"foo");

  XCTAssertEqualObjects([raw.itemArray[2] message], DataFromCStr("baz"));
}

- (void)testParseMessageSet {
  // Set up a MSetRawMessageSet with two known messages and an unknown one.
  MSetRawMessageSet* raw = [MSetRawMessageSet message];

  {
    MSetRawMessageSet_Item* item = [MSetRawMessageSet_Item message];
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
    item.typeId = MSetMessageExtension1_extension_MessageSetExtension().fieldNumber;
#else
    item.typeId = [MSetMessageExtension1 messageSetExtension].fieldNumber;
#endif
    MSetMessageExtension1* message = [MSetMessageExtension1 message];
    message.i = 123;
    item.message = [message data];
    [raw.itemArray addObject:item];
  }

  {
    MSetRawMessageSet_Item* item = [MSetRawMessageSet_Item message];
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
    item.typeId = MSetMessageExtension2_extension_MessageSetExtension().fieldNumber;
#else
    item.typeId = [MSetMessageExtension2 messageSetExtension].fieldNumber;
#endif
    MSetMessageExtension2* message = [MSetMessageExtension2 message];
    message.str = @"foo";
    item.message = [message data];
    [raw.itemArray addObject:item];
  }

  {
    MSetRawMessageSet_Item* item = [MSetRawMessageSet_Item message];
    item.typeId = kUnknownTypeId;
    item.message = DataFromCStr("bar");
    [raw.itemArray addObject:item];
  }

  NSData* data = [raw data];

  // Parse as a MSetMessage and check the contents.
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  MSetMessage* messageSet =
      [MSetMessage parseFromData:data
               extensionRegistry:MSet_Objc_Protobuf_Tests_Mset_MSetUnittestMsetRoot_Registry()
                           error:NULL];

  XCTAssertEqual(
      [[messageSet getExtension:MSetMessageExtension1_extension_MessageSetExtension()] i], 123);
  XCTAssertEqualObjects(
      [[messageSet getExtension:MSetMessageExtension2_extension_MessageSetExtension()] str],
      @"foo");
#else
  MSetMessage* messageSet = [MSetMessage parseFromData:data
                                     extensionRegistry:[MSetUnittestMsetRoot extensionRegistry]
                                                 error:NULL];

  XCTAssertEqual([[messageSet getExtension:[MSetMessageExtension1 messageSetExtension]] i], 123);
  XCTAssertEqualObjects([[messageSet getExtension:[MSetMessageExtension2 messageSetExtension]] str],
                        @"foo");
#endif

  GPBUnknownFields* ufs = [[[GPBUnknownFields alloc] initFromMessage:messageSet] autorelease];
  XCTAssertEqual(ufs.count, (NSUInteger)1);
  GPBUnknownFields* group = [ufs firstGroup:GPBWireFormatMessageSetItem];
  XCTAssertNotNil(group);
  XCTAssertEqual(group.count, (NSUInteger)2);
  uint64_t varint = 0;
  XCTAssertTrue([group getFirst:GPBWireFormatMessageSetTypeId varint:&varint]);
  XCTAssertEqual(varint, kUnknownTypeId);
  XCTAssertEqualObjects([group firstLengthDelimited:GPBWireFormatMessageSetMessage],
                        DataFromCStr("bar"));
}

- (void)testParseMessageSet_FirstValueSticks {
  MSetRawBreakableMessageSet* raw = [MSetRawBreakableMessageSet message];

  {
    MSetRawBreakableMessageSet_Item* item = [MSetRawBreakableMessageSet_Item message];

#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
    [item.typeIdArray addValue:MSetMessageExtension1_extension_MessageSetExtension().fieldNumber];
#else
    [item.typeIdArray addValue:[MSetMessageExtension1 messageSetExtension].fieldNumber];
#endif
    MSetMessageExtension1* message1 = [MSetMessageExtension1 message];
    message1.i = 123;
    NSData* itemData = [message1 data];
    [item.messageArray addObject:itemData];

#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
    [item.typeIdArray addValue:MSetMessageExtension2_extension_MessageSetExtension().fieldNumber];
#else
    [item.typeIdArray addValue:[MSetMessageExtension2 messageSetExtension].fieldNumber];
#endif
    MSetMessageExtension2* message2 = [MSetMessageExtension2 message];
    message2.str = @"foo";
    itemData = [message2 data];
    [item.messageArray addObject:itemData];

    [raw.itemArray addObject:item];
  }

  NSData* data = [raw data];

  // Parse as a MSetMessage and check the contents.
  NSError* err = nil;
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  MSetMessage* messageSet =
      [MSetMessage parseFromData:data
               extensionRegistry:MSet_Objc_Protobuf_Tests_Mset_MSetUnittestMsetRoot_Registry()
                           error:&err];
  XCTAssertNotNil(messageSet);
  XCTAssertNil(err);
  XCTAssertTrue([messageSet hasExtension:MSetMessageExtension1_extension_MessageSetExtension()]);
  XCTAssertEqual(
      [[messageSet getExtension:MSetMessageExtension1_extension_MessageSetExtension()] i], 123);
  XCTAssertFalse([messageSet hasExtension:MSetMessageExtension2_extension_MessageSetExtension()]);
#else
  MSetMessage* messageSet = [MSetMessage parseFromData:data
                                     extensionRegistry:[MSetUnittestMsetRoot extensionRegistry]
                                                 error:&err];
  XCTAssertNotNil(messageSet);
  XCTAssertNil(err);
  XCTAssertTrue([messageSet hasExtension:[MSetMessageExtension1 messageSetExtension]]);
  XCTAssertEqual([[messageSet getExtension:[MSetMessageExtension1 messageSetExtension]] i], 123);
  XCTAssertFalse([messageSet hasExtension:[MSetMessageExtension2 messageSetExtension]]);
#endif
  GPBUnknownFields* ufs = [[[GPBUnknownFields alloc] initFromMessage:messageSet] autorelease];
  XCTAssertTrue(ufs.empty);
}

- (void)testParseMessageSet_PartialValuesDropped {
  MSetRawBreakableMessageSet* raw = [MSetRawBreakableMessageSet message];

  {
    MSetRawBreakableMessageSet_Item* item = [MSetRawBreakableMessageSet_Item message];
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
    [item.typeIdArray addValue:MSetMessageExtension1_extension_MessageSetExtension().fieldNumber];
#else
    [item.typeIdArray addValue:[MSetMessageExtension1 messageSetExtension].fieldNumber];
#endif
    // No payload.
    [raw.itemArray addObject:item];
  }

  {
    MSetRawBreakableMessageSet_Item* item = [MSetRawBreakableMessageSet_Item message];
    // No type ID.
    MSetMessageExtension2* message = [MSetMessageExtension2 message];
    message.str = @"foo";
    NSData* itemData = [message data];
    [item.messageArray addObject:itemData];
    [raw.itemArray addObject:item];
  }

  {
    MSetRawBreakableMessageSet_Item* item = [MSetRawBreakableMessageSet_Item message];
    // Neither type ID nor payload.
    [raw.itemArray addObject:item];
  }

  NSData* data = [raw data];

  // Parse as a MSetMessage and check the contents.
  NSError* err = nil;
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  MSetMessage* messageSet =
      [MSetMessage parseFromData:data
               extensionRegistry:MSet_Objc_Protobuf_Tests_Mset_MSetUnittestMsetRoot_Registry()
                           error:&err];
#else
  MSetMessage* messageSet = [MSetMessage parseFromData:data
                                     extensionRegistry:[MSetUnittestMsetRoot extensionRegistry]
                                                 error:&err];
#endif
  XCTAssertNotNil(messageSet);
  XCTAssertNil(err);
  XCTAssertEqual([messageSet extensionsCurrentlySet].count,
                 (NSUInteger)0);  // None because they were all partial and dropped.
  GPBUnknownFields* ufs = [[[GPBUnknownFields alloc] initFromMessage:messageSet] autorelease];
  XCTAssertTrue(ufs.empty);
}

- (void)testMessageSet_normalEncoding {
  MSetMessageEx* msgEx = [MSetMessageEx message];
  MSetMessageExtension1* message1 = [MSetMessageExtension1 message];
  message1.i = 123;
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  [msgEx setExtension:MSetMessageExtension1_extension_DoppelgangerMessageSetExtension()
                value:message1];
#else
  [msgEx setExtension:[MSetMessageExtension1 doppelgangerMessageSetExtension] value:message1];
#endif
  MSetMessageExtension3* message3 = [MSetMessageExtension3 message];
  message3.x = 10;
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  [msgEx setExtension:MSetMessageExtension3_extension_DoppelgangerMessageSetExtension()
                value:message3];
#else
  [msgEx setExtension:[MSetMessageExtension3 doppelgangerMessageSetExtension] value:message3];
#endif

  NSData* data = [msgEx data];
  XCTAssertNotNil(data);

  NSError* err = nil;
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  MSetMessage* msg =
      [MSetMessage parseFromData:data
               extensionRegistry:MSet_Objc_Protobuf_Tests_Mset_MSetUnittestMsetRoot_Registry()
                           error:&err];
  XCTAssertNil(err);
  XCTAssertNotNil(msg);
  XCTAssertEqual([[msg getExtension:MSetMessageExtension1_extension_MessageSetExtension()] i], 123);
#else
  MSetMessage* msg = [MSetMessage parseFromData:data
                              extensionRegistry:[MSetUnittestMsetRoot extensionRegistry]
                                          error:&err];
  XCTAssertNil(err);
  XCTAssertNotNil(msg);
  XCTAssertEqual([[msg getExtension:[MSetMessageExtension1 messageSetExtension]] i], 123);
#endif
  // Extension 3 is unknown on the actually test MessageSet, so it will stay in unknown fields
  // without being transformed into the group structure.
  GPBUnknownFields* ufs = [[[GPBUnknownFields alloc] initFromMessage:msg] autorelease];
  XCTAssertEqual(ufs.count, (NSUInteger)1);
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  NSData* bytes =
      [ufs firstLengthDelimited:MSetMessageExtension3_extension_DoppelgangerMessageSetExtension()
                                    .fieldNumber];
#else
  NSData* bytes = [ufs
      firstLengthDelimited:[MSetMessageExtension3 doppelgangerMessageSetExtension].fieldNumber];
#endif
  XCTAssertNotNil(bytes);
  XCTAssertEqualObjects(bytes, [message3 data]);
}

- (void)assertFieldsInOrder:(NSData*)data {
  GPBCodedInputStream* input = [GPBCodedInputStream streamWithData:data];
  int32_t previousTag = 0;

  while (YES) {
    int32_t tag = [input readTag];
    if (tag == 0) {
      break;
    }

    XCTAssertGreaterThan(tag, previousTag);
    [input skipField:tag];
  }
}

static NSData* MessageSetDataWithLayers(NSUInteger layers) {
  MSetMessage* innermost = [MSetMessage message];
  MSetMessageExtension1* innermostExt = [MSetMessageExtension1 message];
  innermostExt.i = 1;
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  [innermost setExtension:MSetMessageExtension1_extension_MessageSetExtension() value:innermostExt];
#else
  [innermost setExtension:[MSetMessageExtension1 messageSetExtension] value:innermostExt];
#endif

  MSetMessage* current = innermost;
  for (NSUInteger i = 1; i < layers; ++i) {
    MSetMessageExtension1* ext = [MSetMessageExtension1 message];
    ext.recursive = current;
    MSetMessage* parent = [MSetMessage message];
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
    [parent setExtension:MSetMessageExtension1_extension_MessageSetExtension() value:ext];
#else
    [parent setExtension:[MSetMessageExtension1 messageSetExtension] value:ext];
#endif
    current = parent;
  }
  return [current data];
}

- (void)testParseMessageSetRecursionDepthCarriedFromParent {
  // Each MSetMessage carries a single MSetMessageExtension1, whose
  // `recursive` field is again a MSetMessage. Chaining N of these produces
  // a MessageSet-of-MessageSet payload nested N levels deep. The parser
  // for each MessageSet item allocates a fresh CodedInputStream, so depth
  // tracking has to be inherited across those streams for the documented
  // kDefaultRecursionLimit (100) to actually apply.
  //
  // Each layer increases recursion depth by 2 (+1 for the child CodedInputStream
  // in parseMessageSet:, +1 for the `recursive` message field in readMessage:).
  // 50 layers reaches depth 100 (kDefaultRecursionLimit), which must parse successfully.
  // 51 layers attempts to reach depth 101 (kDefaultRecursionLimit + 1), which must fail
  // with GPBCodedInputStreamErrorRecursionDepthExceeded rather than silently parsing.
  const NSUInteger kPassLayers = 50;
  NSData* passData = MessageSetDataWithLayers(kPassLayers);
  XCTAssertNotNil(passData);

  NSError* error = nil;
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  MSetMessage* passParsed =
      [MSetMessage parseFromData:passData
               extensionRegistry:MSet_Objc_Protobuf_Tests_Mset_MSetUnittestMsetRoot_Registry()
                           error:&error];
#else
  MSetMessage* passParsed = [MSetMessage parseFromData:passData
                                     extensionRegistry:[MSetUnittestMsetRoot extensionRegistry]
                                                 error:&error];
#endif
  XCTAssertNotNil(passParsed);
  XCTAssertNil(error);

  const NSUInteger kFailLayers = 51;
  NSData* failData = MessageSetDataWithLayers(kFailLayers);
  XCTAssertNotNil(failData);

  error = nil;
#if defined(GPB_UNITTEST_USE_C_FUNCTION_FOR_EXTENSIONS)
  MSetMessage* failParsed =
      [MSetMessage parseFromData:failData
               extensionRegistry:MSet_Objc_Protobuf_Tests_Mset_MSetUnittestMsetRoot_Registry()
                           error:&error];
#else
  MSetMessage* failParsed = [MSetMessage parseFromData:failData
                                     extensionRegistry:[MSetUnittestMsetRoot extensionRegistry]
                                                 error:&error];
#endif
  XCTAssertNil(failParsed);
  XCTAssertNotNil(error);
  XCTAssertEqualObjects(error.domain, GPBCodedInputStreamErrorDomain);
  XCTAssertEqual(error.code, GPBCodedInputStreamErrorRecursionDepthExceeded);
}

@end