Fix out-of-memory checking, a compiler warning, and allow subclassing

GCD/GCDAsyncSocket.m

@@ -343,7 +343,7 @@ @interface GCDAsyncSocketPreBuffer : NSObject
 
 - (id)initWithCapacity:(size_t)numBytes;
 
-- (void)ensureCapacityForWrite:(size_t)numBytes;
+- (BOOL)ensureCapacityForWrite:(size_t)numBytes;
 
 - (size_t)availableBytes;
 - (uint8_t *)readBuffer;
@@ -370,7 +370,9 @@ - (id)initWithCapacity:(size_t)numBytes
 	{
 		preBufferSize = numBytes;
 		preBuffer = malloc(preBufferSize);
-
+		if (!preBuffer)
+			return nil;
+
 		readPointer = preBuffer;
 		writePointer = preBuffer;
 	}
@@ -383,8 +385,9 @@ - (void)dealloc
 		free(preBuffer);
 }
 
-- (void)ensureCapacityForWrite:(size_t)numBytes
+- (BOOL)ensureCapacityForWrite:(size_t)numBytes
 {
+	BOOL success;
 	size_t availableSpace = preBufferSize - (writePointer - readPointer);
 
 	if (numBytes > availableSpace)
@@ -393,16 +396,26 @@ - (void)ensureCapacityForWrite:(size_t)numBytes
 
 		size_t newPreBufferSize = preBufferSize + additionalBytes;
 		uint8_t *newPreBuffer = realloc(preBuffer, newPreBufferSize);
-
-		size_t readPointerOffset = readPointer - preBuffer;
-		size_t writePointerOffset = writePointer - preBuffer;
-
-		preBuffer = newPreBuffer;
-		preBufferSize = newPreBufferSize;
-
-		readPointer = preBuffer + readPointerOffset;
-		writePointer = preBuffer + writePointerOffset;
+		if (newPreBuffer)
+		{
+			size_t readPointerOffset = readPointer - preBuffer;
+			size_t writePointerOffset = writePointer - preBuffer;
+
+			preBuffer = newPreBuffer;
+			preBufferSize = newPreBufferSize;
+
+			readPointer = preBuffer + readPointerOffset;
+			writePointer = preBuffer + writePointerOffset;
+
+			success = YES;
+		}
+		else
+			success = NO;
 	}
+	else
+		success = YES;
+
+	return success;
 }
 
 - (size_t)availableBytes
@@ -1794,9 +1807,9 @@ - (BOOL)doAccept:(int)parentSocketFD
 
 			// Create GCDAsyncSocket instance for accepted socket
 
-			GCDAsyncSocket *acceptedSocket = [[GCDAsyncSocket alloc] initWithDelegate:theDelegate
-			                                                            delegateQueue:delegateQueue
-			                                                              socketQueue:childSocketQueue];
+			GCDAsyncSocket *acceptedSocket = [[[self class] alloc] initWithDelegate:theDelegate
+																	  delegateQueue:delegateQueue
+																		socketQueue:childSocketQueue];
 
 			if (isIPv4)
 				acceptedSocket->socket4FD = childSocketFD;
@@ -2997,6 +3010,17 @@ - (NSError *)connectionClosedError
 	return [NSError errorWithDomain:GCDAsyncSocketErrorDomain code:GCDAsyncSocketClosedError userInfo:userInfo];
 }
 
+- (NSError *)outOfMemoryError
+{
+	NSString *errMsg = NSLocalizedStringWithDefaultValue(@"GCDAsyncSocketOutOfMemoryError",
+	                                                     @"GCDAsyncSocket", [NSBundle mainBundle],
+	                                                     @"Out of memory", nil);
+
+	NSDictionary *userInfo = [NSDictionary dictionaryWithObject:errMsg forKey:NSLocalizedDescriptionKey];
+
+	return [NSError errorWithDomain:GCDAsyncSocketErrorDomain code:GCDAsyncSocketOtherError userInfo:userInfo];
+}
+
 - (NSError *)otherError:(NSString *)errMsg
 {
 	NSDictionary *userInfo = [NSDictionary dictionaryWithObject:errMsg forKey:NSLocalizedDescriptionKey];
@@ -4077,19 +4101,20 @@ - (void)flushSSLBuffers
 
 			CFIndex defaultBytesToRead = (1024 * 4);
 
-			[preBuffer ensureCapacityForWrite:defaultBytesToRead];
-
-			uint8_t *buffer = [preBuffer writeBuffer];
-
-			CFIndex result = CFReadStreamRead(readStream, buffer, defaultBytesToRead);
-			LogVerbose(@"%@ - CFReadStreamRead(): result = %i", THIS_METHOD, (int)result);
-
-			if (result > 0)
+			if ([preBuffer ensureCapacityForWrite:defaultBytesToRead])
 			{
-				[preBuffer didWrite:result];
+				uint8_t *buffer = [preBuffer writeBuffer];
+
+				CFIndex result = CFReadStreamRead(readStream, buffer, defaultBytesToRead);
+				LogVerbose(@"%@ - CFReadStreamRead(): result = %i", THIS_METHOD, (int)result);
+
+				if (result > 0)
+				{
+					[preBuffer didWrite:result];
+				}
+
+				flags &= ~kSecureSocketHasBytesAvailable;
 			}
-
-			flags &= ~kSecureSocketHasBytesAvailable;
 		}
 
 		return;
@@ -4133,32 +4158,33 @@ - (void)flushSSLBuffers
 
 			// Make sure there's enough room in the prebuffer
 
-			[preBuffer ensureCapacityForWrite:estimatedBytesAvailable];
-
-			// Read data into prebuffer
-
-			uint8_t *buffer = [preBuffer writeBuffer];
-			size_t bytesRead = 0;
-
-			OSStatus result = SSLRead(sslContext, buffer, (size_t)estimatedBytesAvailable, &bytesRead);
-			LogVerbose(@"%@ - read from secure socket = %u", THIS_METHOD, (unsigned)bytesRead);
-
-			if (bytesRead > 0)
-			{
-				[preBuffer didWrite:bytesRead];
-			}
-
-			LogVerbose(@"%@ - prebuffer.length = %zu", THIS_METHOD, [preBuffer availableBytes]);
-
-			if (result != noErr)
+			if ([preBuffer ensureCapacityForWrite:estimatedBytesAvailable])
 			{
+				// Read data into prebuffer
+
+				uint8_t *buffer = [preBuffer writeBuffer];
+				size_t bytesRead = 0;
+
+				OSStatus result = SSLRead(sslContext, buffer, (size_t)estimatedBytesAvailable, &bytesRead);
+				LogVerbose(@"%@ - read from secure socket = %u", THIS_METHOD, (unsigned)bytesRead);
+
+				if (bytesRead > 0)
+				{
+					[preBuffer didWrite:bytesRead];
+				}
+
+				LogVerbose(@"%@ - prebuffer.length = %zu", THIS_METHOD, [preBuffer availableBytes]);
+
+				if (result != noErr)
+				{
+					done = YES;
+				}
+				else
+				{
+					updateEstimatedBytesAvailable();
+				}
+			} else
 				done = YES;
-			}
-			else
-			{
-				updateEstimatedBytesAvailable();
-			}
-
 		} while (!done && estimatedBytesAvailable > 0);
 	}
 
@@ -4233,7 +4259,11 @@ - (void)doReadData
 		else
 			hasBytesAvailable = NO;
 
-		#endif
+		#else
+        // This should never occur because usingCFStreamForTLS is only TRUE on iPhone, but it silences
+        // an uninitialized variable warning.
+            hasBytesAvailable = NO;
+        #endif
 	}
 	else
 	{
@@ -4492,9 +4522,14 @@ - (void)doReadData
 
 		if (readIntoPreBuffer)
 		{
-			[preBuffer ensureCapacityForWrite:bytesToRead];
-
-			buffer = [preBuffer writeBuffer];
+			if ([preBuffer ensureCapacityForWrite:bytesToRead])
+				buffer = [preBuffer writeBuffer];
+			else
+			{
+				error = [self outOfMemoryError];
+				[self closeWithError:error];
+				return;
+			}
 		}
 		else
 		{
@@ -4727,19 +4762,26 @@ - (void)doReadData
 						// Copy excess data into preBuffer
 
 						LogVerbose(@"copying %ld overflow bytes into preBuffer", (long)overflow);
-						[preBuffer ensureCapacityForWrite:overflow];
-
-						uint8_t *overflowBuffer = buffer + underflow;
-						memcpy([preBuffer writeBuffer], overflowBuffer, overflow);
-
-						[preBuffer didWrite:overflow];
-						LogVerbose(@"preBuffer.length = %zu", [preBuffer availableBytes]);
-
-						// Note: The completeCurrentRead method will trim the buffer for us.
-
-						currentRead->bytesDone += underflow;
-						totalBytesReadForCurrentRead += underflow;
-						done = YES;
+						if ([preBuffer ensureCapacityForWrite:overflow])
+						{
+							uint8_t *overflowBuffer = buffer + underflow;
+							memcpy([preBuffer writeBuffer], overflowBuffer, overflow);
+
+							[preBuffer didWrite:overflow];
+							LogVerbose(@"preBuffer.length = %zu", [preBuffer availableBytes]);
+
+							// Note: The completeCurrentRead method will trim the buffer for us.
+
+							currentRead->bytesDone += underflow;
+							totalBytesReadForCurrentRead += underflow;
+							done = YES;
+						}
+						else
+						{
+							error = [self outOfMemoryError];
+							[self closeWithError:error];
+							return;
+						}
 					}
 					else
 					{
@@ -5959,11 +6001,17 @@ - (OSStatus)sslReadWithBuffer:(void *)buffer length:(size_t *)bufferLength
 
 			LogVerbose(@"%@: Reading into sslPreBuffer...", THIS_METHOD);
 
-			[sslPreBuffer ensureCapacityForWrite:socketFDBytesAvailable];
-
-			readIntoPreBuffer = YES;
-			bytesToRead = (size_t)socketFDBytesAvailable;
-			buf = [sslPreBuffer writeBuffer];
+			if ([sslPreBuffer ensureCapacityForWrite:socketFDBytesAvailable])
+			{
+				readIntoPreBuffer = YES;
+				bytesToRead = (size_t)socketFDBytesAvailable;
+				buf = [sslPreBuffer writeBuffer];
+			}
+			else
+			{
+				socketFDBytesAvailable = 0;
+				return errSSLClosedAbort;
+			}
 		}
 		else
 		{
@@ -6481,14 +6529,18 @@ - (void)ssl_startTLS
 
 	if (preBufferLength > 0)
 	{
-		[sslPreBuffer ensureCapacityForWrite:preBufferLength];
-
-		memcpy([sslPreBuffer writeBuffer], [preBuffer readBuffer], preBufferLength);
-		[preBuffer didRead:preBufferLength];
-		[sslPreBuffer didWrite:preBufferLength];
+		if ([sslPreBuffer ensureCapacityForWrite:preBufferLength])
+		{
+			memcpy([sslPreBuffer writeBuffer], [preBuffer readBuffer], preBufferLength);
+			[preBuffer didRead:preBufferLength];
+			[sslPreBuffer didWrite:preBufferLength];
+			sslErrCode = noErr;
+		}
+		else
+			sslErrCode = errSSLClosedAbort;
 	}
-
-	sslErrCode = noErr;
+	else
+		sslErrCode = noErr;
 
 	// Start the SSL Handshake process
 
@@ -6712,11 +6764,14 @@ + (void)cfstreamThread { @autoreleasepool
 
 	// We can't run the run loop unless it has an associated input source or a timer.
 	// So we'll just create a timer that will never fire - unless the server runs for decades.
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wundeclared-selector"
 	[NSTimer scheduledTimerWithTimeInterval:[[NSDate distantFuture] timeIntervalSinceNow]
 	                                 target:self
 	                               selector:@selector(ignore:)
 	                               userInfo:nil
 	                                repeats:YES];
+#pragma clang diagnostic pop
 
 	[[NSRunLoop currentRunLoop] run];
 

GCD/GCDAsyncUdpSocket.m

@@ -424,10 +424,12 @@ - (id)initWithDelegate:(id)aDelegate delegateQueue:(dispatch_queue_t)dq socketQu
 		sendQueue = [[NSMutableArray alloc] initWithCapacity:5];
 
 		#if TARGET_OS_IPHONE
-		[[NSNotificationCenter defaultCenter] addObserver:self
-		                                         selector:@selector(applicationWillEnterForeground:)
-		                                             name:UIApplicationWillEnterForegroundNotification
-		                                           object:nil];
+            #ifndef TARGET_OS_ATV2
+                [[NSNotificationCenter defaultCenter] addObserver:self
+                                                         selector:@selector(applicationWillEnterForeground:)
+                                                             name:UIApplicationWillEnterForegroundNotification
+                                                           object:nil];
+            #endif
 		#endif
 	}
 	return self;