Skip to content

Commit e6d01b2

Browse files
[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
1 parent c8e005d commit e6d01b2

File tree

5 files changed

+93
-12
lines changed

5 files changed

+93
-12
lines changed

objectivec/GPBCodedOutputStream.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,11 @@ __attribute__((objc_subclassing_restricted))
109109
**/
110110
- (void)flush;
111111

112+
/**
113+
* @return The number of bytes written out. Includes bytes not yet flused.
114+
**/
115+
- (size_t)bytesWritten;
116+
112117
/**
113118
* Write the raw byte out.
114119
*

objectivec/GPBCodedOutputStream.m

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
uint8_t *bytes;
4949
size_t size;
5050
size_t position;
51+
size_t bytesFlushed;
5152
NSOutputStream *output;
5253
} GPBOutputBufferState;
5354

@@ -71,6 +72,7 @@ static void GPBRefreshBuffer(GPBOutputBufferState *state) {
7172
if (written != (NSInteger)state->position) {
7273
[NSException raise:GPBCodedOutputStreamException_WriteFailed format:@""];
7374
}
75+
state->bytesFlushed += written;
7476
state->position = 0;
7577
}
7678
}
@@ -192,6 +194,13 @@ + (instancetype)streamWithData:(NSMutableData *)data {
192194
#pragma clang diagnostic push
193195
#pragma clang diagnostic ignored "-Wdirect-ivar-access"
194196

197+
- (size_t)bytesWritten {
198+
// Could use NSStreamFileCurrentOffsetKey on state_.output if there is a stream, that could be
199+
// expensive, manually tracking what is flush keeps things faster so message serialization can
200+
// check it.
201+
return state_.bytesFlushed + state_.position;
202+
}
203+
195204
- (void)writeDoubleNoTag:(double)value {
196205
GPBWriteRawLittleEndian64(&state_, GPBConvertDoubleToInt64(value));
197206
}
@@ -886,6 +895,7 @@ - (void)writeRawPtr:(const void *)value offset:(size_t)offset length:(size_t)len
886895
if (written != (NSInteger)length) {
887896
[NSException raise:GPBCodedOutputStreamException_WriteFailed format:@""];
888897
}
898+
state_.bytesFlushed += written;
889899
}
890900
}
891901
}

objectivec/GPBMessage.h

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,12 @@ typedef NS_ENUM(NSInteger, GPBMessageErrorCode) {
6161
**/
6262
extern NSString *const GPBErrorReasonKey;
6363

64+
/**
65+
* An exception name raised during serialization when the message would be
66+
* larger than the 2GB limit.
67+
**/
68+
extern NSString *const GPBMessageExceptionMessageTooLarge;
69+
6470
CF_EXTERN_C_END
6571

6672
/**

objectivec/GPBMessage.m

Lines changed: 52 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,16 @@
5959

6060
static NSString *const kGPBDataCoderKey = @"GPBData";
6161

62+
// Length-delimited has a max size of 2GB, and thus messages do also.
63+
// src/google/protobuf/message_lite also does this enforcement on the C++ side. Validation for
64+
// parsing is done with GPBCodedInputStream; but for messages, it is less checks to do it within
65+
// the message side since the input stream code calls these same bottlenecks.
66+
// https://protobuf.dev/programming-guides/encoding/#cheat-sheet
67+
static const size_t kMaximumMessageSize = 0x7fffffff;
68+
69+
NSString *const GPBMessageExceptionMessageTooLarge =
70+
GPBNSStringifySymbol(GPBMessageExceptionMessageTooLarge);
71+
6272
//
6373
// PLEASE REMEMBER:
6474
//
@@ -111,7 +121,7 @@ static id CreateMapForField(GPBFieldDescriptor *field, GPBMessage *autocreator)
111121
__attribute__((ns_returns_retained));
112122
static GPBUnknownFieldSet *GetOrMakeUnknownFields(GPBMessage *self);
113123

114-
#ifdef DEBUG
124+
#if defined(DEBUG) && DEBUG
115125
static NSError *MessageError(NSInteger code, NSDictionary *userInfo) {
116126
return [NSError errorWithDomain:GPBMessageErrorDomain code:code userInfo:userInfo];
117127
}
@@ -968,7 +978,7 @@ - (instancetype)initWithData:(NSData *)data
968978
if (![self mergeFromData:data extensionRegistry:extensionRegistry error:errorPtr]) {
969979
[self release];
970980
self = nil;
971-
#ifdef DEBUG
981+
#if defined(DEBUG) && DEBUG
972982
} else if (!self.initialized) {
973983
[self release];
974984
self = nil;
@@ -997,7 +1007,7 @@ - (instancetype)initWithCodedInputStream:(GPBCodedInputStream *)input
9971007
*errorPtr = ErrorFromException(exception);
9981008
}
9991009
}
1000-
#ifdef DEBUG
1010+
#if defined(DEBUG) && DEBUG
10011011
if (self && !self.initialized) {
10021012
[self release];
10031013
self = nil;
@@ -1290,12 +1300,16 @@ - (GPBDescriptor *)descriptor {
12901300
}
12911301

12921302
- (NSData *)data {
1293-
#ifdef DEBUG
1303+
#if defined(DEBUG) && DEBUG
12941304
if (!self.initialized) {
12951305
return nil;
12961306
}
12971307
#endif
1298-
NSMutableData *data = [NSMutableData dataWithLength:[self serializedSize]];
1308+
size_t expectedSize = [self serializedSize];
1309+
if (expectedSize > kMaximumMessageSize) {
1310+
return nil;
1311+
}
1312+
NSMutableData *data = [NSMutableData dataWithLength:expectedSize];
12991313
GPBCodedOutputStream *stream = [[GPBCodedOutputStream alloc] initWithData:data];
13001314
@try {
13011315
[self writeToCodedOutputStream:stream];
@@ -1306,11 +1320,14 @@ - (NSData *)data {
13061320
// to this message or a message used as a nested field, and that other thread mutated that
13071321
// message, causing the pre computed serializedSize to no longer match the final size after
13081322
// serialization. It is not safe to mutate a message while accessing it from another thread.
1309-
#ifdef DEBUG
1323+
#if defined(DEBUG) && DEBUG
13101324
NSLog(@"%@: Internal exception while building message data: %@", [self class], exception);
13111325
#endif
13121326
data = nil;
13131327
}
1328+
#if defined(DEBUG) && DEBUG
1329+
NSAssert(!data || [stream bytesWritten] == expectedSize, @"Internal error within the library");
1330+
#endif
13141331
[stream release];
13151332
return data;
13161333
}
@@ -1329,7 +1346,7 @@ - (NSData *)delimitedData {
13291346
// to this message or a message used as a nested field, and that other thread mutated that
13301347
// message, causing the pre computed serializedSize to no longer match the final size after
13311348
// serialization. It is not safe to mutate a message while accessing it from another thread.
1332-
#ifdef DEBUG
1349+
#if defined(DEBUG) && DEBUG
13331350
NSLog(@"%@: Internal exception while building message delimitedData: %@", [self class],
13341351
exception);
13351352
#endif
@@ -1342,8 +1359,16 @@ - (NSData *)delimitedData {
13421359

13431360
- (void)writeToOutputStream:(NSOutputStream *)output {
13441361
GPBCodedOutputStream *stream = [[GPBCodedOutputStream alloc] initWithOutputStream:output];
1345-
[self writeToCodedOutputStream:stream];
1346-
[stream release];
1362+
@try {
1363+
[self writeToCodedOutputStream:stream];
1364+
size_t bytesWritten = [stream bytesWritten];
1365+
if (bytesWritten > kMaximumMessageSize) {
1366+
[NSException raise:GPBMessageExceptionMessageTooLarge
1367+
format:@"Message would have been %zu bytes", bytesWritten];
1368+
}
1369+
} @finally {
1370+
[stream release];
1371+
}
13471372
}
13481373

13491374
- (void)writeToCodedOutputStream:(GPBCodedOutputStream *)output {
@@ -1377,13 +1402,28 @@ - (void)writeToCodedOutputStream:(GPBCodedOutputStream *)output {
13771402

13781403
- (void)writeDelimitedToOutputStream:(NSOutputStream *)output {
13791404
GPBCodedOutputStream *codedOutput = [[GPBCodedOutputStream alloc] initWithOutputStream:output];
1380-
[self writeDelimitedToCodedOutputStream:codedOutput];
1381-
[codedOutput release];
1405+
@try {
1406+
[self writeDelimitedToCodedOutputStream:codedOutput];
1407+
} @finally {
1408+
[codedOutput release];
1409+
}
13821410
}
13831411

13841412
- (void)writeDelimitedToCodedOutputStream:(GPBCodedOutputStream *)output {
1385-
[output writeRawVarintSizeTAs32:[self serializedSize]];
1413+
size_t expectedSize = [self serializedSize];
1414+
if (expectedSize > kMaximumMessageSize) {
1415+
[NSException raise:GPBMessageExceptionMessageTooLarge
1416+
format:@"Message would have been %zu bytes", expectedSize];
1417+
}
1418+
[output writeRawVarintSizeTAs32:expectedSize];
1419+
#if defined(DEBUG) && DEBUG && !defined(NS_BLOCK_ASSERTIONS)
1420+
size_t initialSize = [output bytesWritten];
1421+
#endif
13861422
[self writeToCodedOutputStream:output];
1423+
#if defined(DEBUG) && DEBUG && !defined(NS_BLOCK_ASSERTIONS)
1424+
NSAssert(([output bytesWritten] - initialSize) == expectedSize,
1425+
@"Internal error within the library");
1426+
#endif
13871427
}
13881428

13891429
- (void)writeField:(GPBFieldDescriptor *)field toCodedOutputStream:(GPBCodedOutputStream *)output {

objectivec/Tests/GPBCodedOutputStreamTests.m

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,9 @@ - (void)assertWriteLittleEndian32:(NSData*)data value:(int32_t)value {
8080
NSOutputStream* rawOutput = [NSOutputStream outputStreamToMemory];
8181
GPBCodedOutputStream* output = [GPBCodedOutputStream streamWithOutputStream:rawOutput];
8282
[output writeRawLittleEndian32:(int32_t)value];
83+
XCTAssertEqual(output.bytesWritten, data.length);
8384
[output flush];
85+
XCTAssertEqual(output.bytesWritten, data.length);
8486

8587
NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
8688
XCTAssertEqualObjects(data, actual);
@@ -90,7 +92,9 @@ - (void)assertWriteLittleEndian32:(NSData*)data value:(int32_t)value {
9092
rawOutput = [NSOutputStream outputStreamToMemory];
9193
output = [GPBCodedOutputStream streamWithOutputStream:rawOutput bufferSize:blockSize];
9294
[output writeRawLittleEndian32:(int32_t)value];
95+
XCTAssertEqual(output.bytesWritten, data.length);
9396
[output flush];
97+
XCTAssertEqual(output.bytesWritten, data.length);
9498

9599
actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
96100
XCTAssertEqualObjects(data, actual);
@@ -101,7 +105,9 @@ - (void)assertWriteLittleEndian64:(NSData*)data value:(int64_t)value {
101105
NSOutputStream* rawOutput = [NSOutputStream outputStreamToMemory];
102106
GPBCodedOutputStream* output = [GPBCodedOutputStream streamWithOutputStream:rawOutput];
103107
[output writeRawLittleEndian64:value];
108+
XCTAssertEqual(output.bytesWritten, data.length);
104109
[output flush];
110+
XCTAssertEqual(output.bytesWritten, data.length);
105111

106112
NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
107113
XCTAssertEqualObjects(data, actual);
@@ -111,7 +117,9 @@ - (void)assertWriteLittleEndian64:(NSData*)data value:(int64_t)value {
111117
rawOutput = [NSOutputStream outputStreamToMemory];
112118
output = [GPBCodedOutputStream streamWithOutputStream:rawOutput bufferSize:blockSize];
113119
[output writeRawLittleEndian64:value];
120+
XCTAssertEqual(output.bytesWritten, data.length);
114121
[output flush];
122+
XCTAssertEqual(output.bytesWritten, data.length);
115123

116124
actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
117125
XCTAssertEqualObjects(data, actual);
@@ -124,7 +132,9 @@ - (void)assertWriteVarint:(NSData*)data value:(int64_t)value {
124132
NSOutputStream* rawOutput = [NSOutputStream outputStreamToMemory];
125133
GPBCodedOutputStream* output = [GPBCodedOutputStream streamWithOutputStream:rawOutput];
126134
[output writeRawVarint32:(int32_t)value];
135+
XCTAssertEqual(output.bytesWritten, data.length);
127136
[output flush];
137+
XCTAssertEqual(output.bytesWritten, data.length);
128138

129139
NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
130140
XCTAssertEqualObjects(data, actual);
@@ -137,7 +147,9 @@ - (void)assertWriteVarint:(NSData*)data value:(int64_t)value {
137147
NSOutputStream* rawOutput = [NSOutputStream outputStreamToMemory];
138148
GPBCodedOutputStream* output = [GPBCodedOutputStream streamWithOutputStream:rawOutput];
139149
[output writeRawVarint64:value];
150+
XCTAssertEqual(output.bytesWritten, data.length);
140151
[output flush];
152+
XCTAssertEqual(output.bytesWritten, data.length);
141153

142154
NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
143155
XCTAssertEqualObjects(data, actual);
@@ -155,7 +167,9 @@ - (void)assertWriteVarint:(NSData*)data value:(int64_t)value {
155167
bufferSize:blockSize];
156168

157169
[output writeRawVarint32:(int32_t)value];
170+
XCTAssertEqual(output.bytesWritten, data.length);
158171
[output flush];
172+
XCTAssertEqual(output.bytesWritten, data.length);
159173

160174
NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
161175
XCTAssertEqualObjects(data, actual);
@@ -167,7 +181,9 @@ - (void)assertWriteVarint:(NSData*)data value:(int64_t)value {
167181
bufferSize:blockSize];
168182

169183
[output writeRawVarint64:value];
184+
XCTAssertEqual(output.bytesWritten, data.length);
170185
[output flush];
186+
XCTAssertEqual(output.bytesWritten, data.length);
171187

172188
NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
173189
XCTAssertEqualObjects(data, actual);
@@ -181,7 +197,9 @@ - (void)assertWriteStringNoTag:(NSData*)data
181197
NSOutputStream* rawOutput = [NSOutputStream outputStreamToMemory];
182198
GPBCodedOutputStream* output = [GPBCodedOutputStream streamWithOutputStream:rawOutput];
183199
[output writeStringNoTag:value];
200+
XCTAssertEqual(output.bytesWritten, data.length);
184201
[output flush];
202+
XCTAssertEqual(output.bytesWritten, data.length);
185203

186204
NSData* actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
187205
XCTAssertEqualObjects(data, actual, @"%@", contextMessage);
@@ -191,7 +209,9 @@ - (void)assertWriteStringNoTag:(NSData*)data
191209
rawOutput = [NSOutputStream outputStreamToMemory];
192210
output = [GPBCodedOutputStream streamWithOutputStream:rawOutput bufferSize:blockSize];
193211
[output writeStringNoTag:value];
212+
XCTAssertEqual(output.bytesWritten, data.length);
194213
[output flush];
214+
XCTAssertEqual(output.bytesWritten, data.length);
195215

196216
actual = [rawOutput propertyForKey:NSStreamDataWrittenToMemoryStreamKey];
197217
XCTAssertEqualObjects(data, actual, @"%@", contextMessage);

0 commit comments

Comments
 (0)