[ObjC] Enforce the max message size when serializing to binary form.

The validation is done at the highest point so if a sub message is what
goes over the limit it is caught at the outer message, thus reducing the
impact on the serialization code.

PiperOrigin-RevId: 511473008

objectivec/GPBCodedOutputStream.h

@@ -109,6 +109,11 @@ __attribute__((objc_subclassing_restricted))
  **/
 - (void)flush;
 
+/**
+ * @return The number of bytes written out. Includes bytes not yet flused.
+ **/
+- (size_t)bytesWritten;
+
 /**
  * Write the raw byte out.
  *

objectivec/GPBCodedOutputStream.m

@@ -48,6 +48,7 @@ typedef struct GPBOutputBufferState {
   uint8_t *bytes;
   size_t size;
   size_t position;
+  size_t bytesFlushed;
   NSOutputStream *output;
 } GPBOutputBufferState;
 
@@ -71,6 +72,7 @@ static void GPBRefreshBuffer(GPBOutputBufferState *state) {
     if (written != (NSInteger)state->position) {
       [NSException raise:GPBCodedOutputStreamException_WriteFailed format:@""];
     }
+    state->bytesFlushed += written;
     state->position = 0;
   }
 }
@@ -192,6 +194,13 @@ static void GPBWriteRawLittleEndian64(GPBOutputBufferState *state, int64_t value
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wdirect-ivar-access"
 
+- (size_t)bytesWritten {
+  // Could use NSStreamFileCurrentOffsetKey on state_.output if there is a stream, that could be
+  // expensive, manually tracking what is flush keeps things faster so message serialization can
+  // check it.
+  return state_.bytesFlushed + state_.position;
+}
+
 - (void)writeDoubleNoTag:(double)value {
   GPBWriteRawLittleEndian64(&state_, GPBConvertDoubleToInt64(value));
 }
@@ -886,6 +895,7 @@ static void GPBWriteRawLittleEndian64(GPBOutputBufferState *state, int64_t value
       if (written != (NSInteger)length) {
         [NSException raise:GPBCodedOutputStreamException_WriteFailed format:@""];
       }
+      state_.bytesFlushed += written;
     }
   }
 }

objectivec/GPBMessage.h

@@ -61,6 +61,12 @@ typedef NS_ENUM(NSInteger, GPBMessageErrorCode) {
  **/
 extern NSString *const GPBErrorReasonKey;
 
+/**
+ * An exception name raised during serialization when the message would be
+ * larger than the 2GB limit.
+ **/
+extern NSString *const GPBMessageExceptionMessageTooLarge;
+
 CF_EXTERN_C_END
 
 /**

objectivec/GPBMessage.m

@@ -59,6 +59,16 @@ NSString *const GPBErrorReasonKey = @"Reason";
 
 static NSString *const kGPBDataCoderKey = @"GPBData";
 
+// Length-delimited has a max size of 2GB, and thus messages do also.
+// src/google/protobuf/message_lite also does this enforcement on the C++ side. Validation for
+// parsing is done with GPBCodedInputStream; but for messages, it is less checks to do it within
+// the message side since the input stream code calls these same bottlenecks.
+// https://protobuf.dev/programming-guides/encoding/#cheat-sheet
+static const size_t kMaximumMessageSize = 0x7fffffff;
+
+NSString *const GPBMessageExceptionMessageTooLarge =
+    GPBNSStringifySymbol(GPBMessageExceptionMessageTooLarge);
+
 //
 // PLEASE REMEMBER:
 //
@@ -111,7 +121,7 @@ static NSMutableDictionary *CloneExtensionMap(NSDictionary *extensionMap, NSZone
     __attribute__((ns_returns_retained));
 static GPBUnknownFieldSet *GetOrMakeUnknownFields(GPBMessage *self);
 
-#ifdef DEBUG
+#if defined(DEBUG) && DEBUG
 static NSError *MessageError(NSInteger code, NSDictionary *userInfo) {
   return [NSError errorWithDomain:GPBMessageErrorDomain code:code userInfo:userInfo];
 }
@@ -968,7 +978,7 @@ static GPBUnknownFieldSet *GetOrMakeUnknownFields(GPBMessage *self) {
     if (![self mergeFromData:data extensionRegistry:extensionRegistry error:errorPtr]) {
       [self release];
       self = nil;
-#ifdef DEBUG
+#if defined(DEBUG) && DEBUG
     } else if (!self.initialized) {
       [self release];
       self = nil;
@@ -997,7 +1007,7 @@ static GPBUnknownFieldSet *GetOrMakeUnknownFields(GPBMessage *self) {
         *errorPtr = ErrorFromException(exception);
       }
     }
-#ifdef DEBUG
+#if defined(DEBUG) && DEBUG
     if (self && !self.initialized) {
       [self release];
       self = nil;
@@ -1290,12 +1300,16 @@ static GPBUnknownFieldSet *GetOrMakeUnknownFields(GPBMessage *self) {
 }
 
 - (NSData *)data {
-#ifdef DEBUG
+#if defined(DEBUG) && DEBUG
   if (!self.initialized) {
     return nil;
   }
 #endif
-  NSMutableData *data = [NSMutableData dataWithLength:[self serializedSize]];
+  size_t expectedSize = [self serializedSize];
+  if (expectedSize > kMaximumMessageSize) {
+    return nil;
+  }
+  NSMutableData *data = [NSMutableData dataWithLength:expectedSize];
   GPBCodedOutputStream *stream = [[GPBCodedOutputStream alloc] initWithData:data];
   @try {
     [self writeToCodedOutputStream:stream];
@@ -1306,11 +1320,14 @@ static GPBUnknownFieldSet *GetOrMakeUnknownFields(GPBMessage *self) {
     // to this message or a message used as a nested field, and that other thread mutated that
     // message, causing the pre computed serializedSize to no longer match the final size after
     // serialization. It is not safe to mutate a message while accessing it from another thread.
-#ifdef DEBUG
+#if defined(DEBUG) && DEBUG
     NSLog(@"%@: Internal exception while building message data: %@", [self class], exception);
 #endif
     data = nil;
   }
+#if defined(DEBUG) && DEBUG
+  NSAssert(!data || [stream bytesWritten] == expectedSize, @"Internal error within the library");
+#endif
   [stream release];
   return data;
 }
@@ -1329,7 +1346,7 @@ static GPBUnknownFieldSet *GetOrMakeUnknownFields(GPBMessage *self) {
     // to this message or a message used as a nested field, and that other thread mutated that
     // message, causing the pre computed serializedSize to no longer match the final size after
     // serialization. It is not safe to mutate a message while accessing it from another thread.
-#ifdef DEBUG
+#if defined(DEBUG) && DEBUG
     NSLog(@"%@: Internal exception while building message delimitedData: %@", [self class],
           exception);
 #endif
@@ -1342,8 +1359,16 @@ static GPBUnknownFieldSet *GetOrMakeUnknownFields(GPBMessage *self) {
 
 - (void)writeToOutputStream:(NSOutputStream *)output {
   GPBCodedOutputStream *stream = [[GPBCodedOutputStream alloc] initWithOutputStream:output];
-  [self writeToCodedOutputStream:stream];
+  @try {
-  [stream release];
+    [self writeToCodedOutputStream:stream];
+    size_t bytesWritten = [stream bytesWritten];
+    if (bytesWritten > kMaximumMessageSize) {
+      [NSException raise:GPBMessageExceptionMessageTooLarge
+                  format:@"Message would have been %zu bytes", bytesWritten];
+    }
+  } @finally {
+    [stream release];
+  }
 }
 
 - (void)writeToCodedOutputStream:(GPBCodedOutputStream *)output {
@@ -1377,13 +1402,28 @@ static GPBUnknownFieldSet *GetOrMakeUnknownFields(GPBMessage *self) {
 
 - (void)writeDelimitedToOutputStream:(NSOutputStream *)output {
   GPBCodedOutputStream *codedOutput = [[GPBCodedOutputStream alloc] initWithOutputStream:output];
-  [self writeDelimitedToCodedOutputStream:codedOutput];
+  @try {
-  [codedOutput release];
+    [self writeDelimitedToCodedOutputStream:codedOutput];
+  } @finally {
+    [codedOutput release];
+  }
 }
 
 - (void)writeDelimitedToCodedOutputStream:(GPBCodedOutputStream *)output {
-  [output writeRawVarintSizeTAs32:[self serializedSize]];
+  size_t expectedSize = [self serializedSize];
+  if (expectedSize > kMaximumMessageSize) {
+    [NSException raise:GPBMessageExceptionMessageTooLarge
+                format:@"Message would have been %zu bytes", expectedSize];
+  }
+  [output writeRawVarintSizeTAs32:expectedSize];
+#if defined(DEBUG) && DEBUG && !defined(NS_BLOCK_ASSERTIONS)
+  size_t initialSize = [output bytesWritten];
+#endif
   [self writeToCodedOutputStream:output];
+#if defined(DEBUG) && DEBUG && !defined(NS_BLOCK_ASSERTIONS)
+  NSAssert(([output bytesWritten] - initialSize) == expectedSize,
+           @"Internal error within the library");
+#endif
 }
 
 - (void)writeField:(GPBFieldDescriptor *)field toCodedOutputStream:(GPBCodedOutputStream *)output {

objectivec/Tests/GPBCodedOutputStreamTests.m

@@ -80,7 +80,9 @@
   NSOutputStream* rawOutput = [NSOutputStream outputStreamToMemory];
   GPBCodedOutputStream* output = [GPBCodedOutputStream streamWithOutputStream:rawOutput];
   [output writeRawLittleEndian32:(int32_t)value];
+  XCTAssertEqual(output.bytesWritten, data.length);
   [output flush];
+  XCTAssertEqual(output.bytesWritten, data.length);
 
   NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
   XCTAssertEqualObjects(data, actual);
@@ -90,7 +92,9 @@
     rawOutput = [NSOutputStream outputStreamToMemory];
     output = [GPBCodedOutputStream streamWithOutputStream:rawOutput bufferSize:blockSize];
     [output writeRawLittleEndian32:(int32_t)value];
+    XCTAssertEqual(output.bytesWritten, data.length);
     [output flush];
+    XCTAssertEqual(output.bytesWritten, data.length);
 
     actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
     XCTAssertEqualObjects(data, actual);
@@ -101,7 +105,9 @@
   NSOutputStream* rawOutput = [NSOutputStream outputStreamToMemory];
   GPBCodedOutputStream* output = [GPBCodedOutputStream streamWithOutputStream:rawOutput];
   [output writeRawLittleEndian64:value];
+  XCTAssertEqual(output.bytesWritten, data.length);
   [output flush];
+  XCTAssertEqual(output.bytesWritten, data.length);
 
   NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
   XCTAssertEqualObjects(data, actual);
@@ -111,7 +117,9 @@
     rawOutput = [NSOutputStream outputStreamToMemory];
     output = [GPBCodedOutputStream streamWithOutputStream:rawOutput bufferSize:blockSize];
     [output writeRawLittleEndian64:value];
+    XCTAssertEqual(output.bytesWritten, data.length);
     [output flush];
+    XCTAssertEqual(output.bytesWritten, data.length);
 
     actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
     XCTAssertEqualObjects(data, actual);
@@ -124,7 +132,9 @@
     NSOutputStream* rawOutput = [NSOutputStream outputStreamToMemory];
     GPBCodedOutputStream* output = [GPBCodedOutputStream streamWithOutputStream:rawOutput];
     [output writeRawVarint32:(int32_t)value];
+    XCTAssertEqual(output.bytesWritten, data.length);
     [output flush];
+    XCTAssertEqual(output.bytesWritten, data.length);
 
     NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
     XCTAssertEqualObjects(data, actual);
@@ -137,7 +147,9 @@
     NSOutputStream* rawOutput = [NSOutputStream outputStreamToMemory];
     GPBCodedOutputStream* output = [GPBCodedOutputStream streamWithOutputStream:rawOutput];
     [output writeRawVarint64:value];
+    XCTAssertEqual(output.bytesWritten, data.length);
     [output flush];
+    XCTAssertEqual(output.bytesWritten, data.length);
 
     NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
     XCTAssertEqualObjects(data, actual);
@@ -155,7 +167,9 @@
                                                                        bufferSize:blockSize];
 
       [output writeRawVarint32:(int32_t)value];
+      XCTAssertEqual(output.bytesWritten, data.length);
       [output flush];
+      XCTAssertEqual(output.bytesWritten, data.length);
 
       NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
       XCTAssertEqualObjects(data, actual);
@@ -167,7 +181,9 @@
                                                                        bufferSize:blockSize];
 
       [output writeRawVarint64:value];
+      XCTAssertEqual(output.bytesWritten, data.length);
       [output flush];
+      XCTAssertEqual(output.bytesWritten, data.length);
 
       NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
       XCTAssertEqualObjects(data, actual);
@@ -181,7 +197,9 @@
   NSOutputStream* rawOutput = [NSOutputStream outputStreamToMemory];
   GPBCodedOutputStream* output = [GPBCodedOutputStream streamWithOutputStream:rawOutput];
   [output writeStringNoTag:value];
+  XCTAssertEqual(output.bytesWritten, data.length);
   [output flush];
+  XCTAssertEqual(output.bytesWritten, data.length);
 
   NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
   XCTAssertEqualObjects(data, actual, @"%@", contextMessage);
@@ -191,7 +209,9 @@
     rawOutput = [NSOutputStream outputStreamToMemory];
     output = [GPBCodedOutputStream streamWithOutputStream:rawOutput bufferSize:blockSize];
     [output writeStringNoTag:value];
+    XCTAssertEqual(output.bytesWritten, data.length);
     [output flush];
+    XCTAssertEqual(output.bytesWritten, data.length);
 
     actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
     XCTAssertEqualObjects(data, actual, @"%@", contextMessage);