Moved response body chunked transfer encoding to GCDWebServerConnection

CGDWebServer/GCDWebServerConnection.m

@@ -44,6 +44,7 @@ typedef void (^WriteBodyCompletionBlock)(BOOL success);
 static NSData* _CRLFData = nil;
 static NSData* _CRLFCRLFData = nil;
 static NSData* _continueData = nil;
+static NSData* _lastChunkData = nil;
 static NSDateFormatter* _dateFormatter = nil;
 static dispatch_queue_t _formatterQueue = NULL;
 
@@ -308,6 +309,26 @@ static inline NSUInteger _ScanHexNumber(const void* bytes, NSUInteger size) {
   NSData* data = [_response performReadData:&error];
   if (data) {
     if (data.length) {
+      if (_response.usesChunkedTransferEncoding) {
+        const char* hexString = [[NSString stringWithFormat:@"%lx", (unsigned long)data.length] UTF8String];
+        size_t hexLength = strlen(hexString);
+        NSData* chunk = [NSMutableData dataWithLength:(hexLength + 2 + data.length + 2)];
+        if (chunk == nil) {
+          LOG_ERROR(@"Failed allocating memory for response body chunk for socket %i: %@", _socket, error);
+          block(NO);
+          return;
+        }
+        char* ptr = (char*)[(NSMutableData*)chunk mutableBytes];
+        bcopy(hexString, ptr, hexLength);
+        ptr += hexLength;
+        *ptr++ = '\r';
+        *ptr++ = '\n';
+        bcopy(data.bytes, ptr, data.length);
+        ptr += data.length;
+        *ptr++ = '\r';
+        *ptr = '\n';
+        data = chunk;
+      }
       [self _writeData:data withCompletionBlock:^(BOOL success) {
         
         if (success) {
@@ -318,7 +339,15 @@ static inline NSUInteger _ScanHexNumber(const void* bytes, NSUInteger size) {
         
       }];
     } else {
-      block(YES);
+      if (_response.usesChunkedTransferEncoding) {
+        [self _writeData:_lastChunkData withCompletionBlock:^(BOOL success) {
+          
+          block(success);
+          
+        }];
+      } else {
+        block(YES);
+      }
     }
   } else {
     LOG_ERROR(@"Failed reading response body for socket %i: %@", _socket, error);
@@ -351,6 +380,9 @@ static inline NSUInteger _ScanHexNumber(const void* bytes, NSUInteger size) {
     CFRelease(message);
     DCHECK(_continueData);
   }
+  if (_lastChunkData == nil) {
+    _lastChunkData = [[NSData alloc] initWithBytes:"0\r\n\r\n" length:5];
+  }
   if (_dateFormatter == nil) {
     DCHECK([NSThread isMainThread]);  // NSDateFormatter should be initialized on main thread
     _dateFormatter = [[NSDateFormatter alloc] init];
@@ -412,6 +444,9 @@ static inline NSUInteger _ScanHexNumber(const void* bytes, NSUInteger size) {
     if (_response.contentLength != NSNotFound) {
       CFHTTPMessageSetHeaderFieldValue(_responseMessage, CFSTR("Content-Length"), (ARC_BRIDGE CFStringRef)[NSString stringWithFormat:@"%lu", (unsigned long)_response.contentLength]);
     }
+    if (_response.usesChunkedTransferEncoding) {
+      CFHTTPMessageSetHeaderFieldValue(_responseMessage, CFSTR("Transfer-Encoding"), CFSTR("chunked"));
+    }
     [_response.additionalHeaders enumerateKeysAndObjectsUsingBlock:^(id key, id obj, BOOL* stop) {
       CFHTTPMessageSetHeaderFieldValue(_responseMessage, (ARC_BRIDGE CFStringRef)key, (ARC_BRIDGE CFStringRef)obj);
     }];

CGDWebServer/GCDWebServerPrivate.h

@@ -127,6 +127,7 @@ extern NSStringEncoding GCDWebServerStringEncodingFromCharset(NSString* charset)
 
 @interface GCDWebServerResponse ()
 @property(nonatomic, readonly) NSDictionary* additionalHeaders;
+@property(nonatomic, readonly) BOOL usesChunkedTransferEncoding;
 - (BOOL)performOpen:(NSError**)error;
 - (NSData*)performReadData:(NSError**)error;
 - (void)performClose;

CGDWebServer/GCDWebServerResponse.h

@@ -39,7 +39,6 @@
 @property(nonatomic) NSInteger statusCode;  // Default is 200
 @property(nonatomic) NSUInteger cacheControlMaxAge;  // Default is 0 seconds i.e. "no-cache"
 @property(nonatomic, getter=isGZipContentEncodingEnabled) BOOL gzipContentEncodingEnabled;  // Default is disabled
-@property(nonatomic, getter=isChunkedTransferEncodingEnabled) BOOL chunkedTransferEncodingEnabled;  // Default is disabled
 + (GCDWebServerResponse*) response;
 - (id)init;
 - (void)setValue:(NSString*)value forAdditionalHeader:(NSString*)header;

CGDWebServer/GCDWebServerResponse.m

@@ -36,9 +36,6 @@
 - (id)initWithResponse:(GCDWebServerResponse*)response reader:(id<GCDWebServerBodyReader>)reader;
 @end
 
-@interface GCDWebServerChunkEncoder : GCDWebServerBodyEncoder
-@end
-
 @interface GCDWebServerGZipEncoder : GCDWebServerBodyEncoder
 @end
 
@@ -73,59 +70,6 @@
 
 @end
 
-@interface GCDWebServerChunkEncoder () {
-@private
-  BOOL _finished;
-}
-@end
-
-@implementation GCDWebServerChunkEncoder
-
-- (id)initWithResponse:(GCDWebServerResponse*)response reader:(id<GCDWebServerBodyReader>)reader {
-  if ((self = [super initWithResponse:response reader:reader])) {
-    response.contentLength = NSNotFound;  // Make sure "Content-Length" header is not set since body length is determined by chunked transfer encoding
-    [response setValue:@"chunked" forAdditionalHeader:@"Transfer-Encoding"];
-  }
-  return self;
-}
-
-- (NSData*)readData:(NSError**)error {
-  NSData* chunk;
-  if (_finished) {
-    chunk = [[NSData alloc] init];
-  } else {
-    NSData* data = [super readData:error];
-    if (data == nil) {
-      return nil;
-    }
-    if (data.length) {
-      const char* hexString = [[NSString stringWithFormat:@"%lx", (unsigned long)data.length] UTF8String];
-      size_t hexLength = strlen(hexString);
-      chunk = [[NSMutableData alloc] initWithLength:(hexLength + 2 + data.length + 2)];
-      if (chunk == nil) {
-        DNOT_REACHED();
-        return nil;
-      }
-      char* ptr = (char*)[(NSMutableData*)chunk mutableBytes];
-      bcopy(hexString, ptr, hexLength);
-      ptr += hexLength;
-      *ptr++ = '\r';
-      *ptr++ = '\n';
-      bcopy(data.bytes, ptr, data.length);
-      ptr += data.length;
-      *ptr++ = '\r';
-      *ptr = '\n';
-    } else {
-      chunk = [[NSData alloc] initWithBytes:"0\r\n\r\n" length:5];
-      DCHECK(chunk);
-      _finished = YES;
-    }
-  }
-  return ARC_AUTORELEASE(chunk);
-}
-
-@end
-
 @interface GCDWebServerGZipEncoder () {
 @private
   z_stream _stream;
@@ -137,7 +81,7 @@
 
 - (id)initWithResponse:(GCDWebServerResponse*)response reader:(id<GCDWebServerBodyReader>)reader {
   if ((self = [super initWithResponse:response reader:reader])) {
-    response.contentLength = NSNotFound;  // Make sure "Content-Length" header is not set since body length is determined by closing the connection
+    response.contentLength = NSNotFound;  // Make sure "Content-Length" header is not set since we don't know it (client will determine body length when connection is closed)
     [response setValue:@"gzip" forAdditionalHeader:@"Content-Encoding"];
   }
   return self;
@@ -213,8 +157,8 @@
   NSInteger _status;
   NSUInteger _maxAge;
   NSMutableDictionary* _headers;
-  BOOL _gzipped;
   BOOL _chunked;
+  BOOL _gzipped;
   
   BOOL _opened;
   NSMutableArray* _encoders;
@@ -225,7 +169,7 @@
 @implementation GCDWebServerResponse
 
 @synthesize contentType=_type, contentLength=_length, statusCode=_status, cacheControlMaxAge=_maxAge,
-            gzipContentEncodingEnabled=_gzipped, chunkedTransferEncodingEnabled=_chunked, additionalHeaders=_headers;
+            gzipContentEncodingEnabled=_gzipped, additionalHeaders=_headers;
 
 + (GCDWebServerResponse*)response {
   return ARC_AUTORELEASE([[[self class] alloc] init]);
@@ -259,6 +203,10 @@
   return _type ? YES : NO;
 }
 
+- (BOOL)usesChunkedTransferEncoding {
+  return (_type != nil) && (_length == NSNotFound);
+}
+
 - (BOOL)open:(NSError**)error {
   return YES;
 }
@@ -286,12 +234,6 @@
     ARC_RELEASE(encoder);
     _reader = encoder;
   }
-  if (_chunked) {
-    GCDWebServerChunkEncoder* encoder = [[GCDWebServerChunkEncoder alloc] initWithResponse:self reader:_reader];
-    [_encoders addObject:encoder];
-    ARC_RELEASE(encoder);
-    _reader = encoder;
-  }
   return [_reader open:error];
 }
 

CGDWebServer/GCDWebServerStreamResponse.m

@@ -44,7 +44,6 @@
     _block = [block copy];
     
     self.contentType = type;
-    self.chunkedTransferEncodingEnabled = YES;
   }
   return self;
 }