Make direct PUT to attachment URL work with streamed request body

This case was accidentally not fixed back in commit 7f24896. Fixes #198.
couchbaselabs · Dec 12, 2012 · d752da8 · d752da8
1 parent 8f6ba32
commit d752da8
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 18 deletions.
diff --git a/Source/TDRouter+Handlers.m b/Source/TDRouter+Handlers.m
@@ -854,7 +854,10 @@ - (TDStatus) do_DELETE: (TD_Database*)db docID: (NSString*)docID {
 }
 
 
-- (TDStatus) updateAttachment: (NSString*)attachment docID: (NSString*)docID body: (NSData*)body {
+- (TDStatus) updateAttachment: (NSString*)attachment
+                        docID: (NSString*)docID
+                         body: (TDBlobStoreWriter*)body
+{
     TDStatus status;
     TD_Revision* rev = [_db updateAttachment: attachment 
                                        body: body
@@ -874,16 +877,35 @@ - (TDStatus) updateAttachment: (NSString*)attachment docID: (NSString*)docID bod
 
 
 - (TDStatus) do_PUT: (TD_Database*)db docID: (NSString*)docID attachment: (NSString*)attachment {
-    return [self updateAttachment: attachment
+    TDBlobStoreWriter* blob = db.attachmentWriter;
-                            docID: docID
+    NSInputStream* bodyStream = _request.HTTPBodyStream;
-                             body: (_request.HTTPBody ?: [NSData data])];
+    if (bodyStream) {
+        // OPT: Should read this asynchronously
+        NSMutableData* buffer = [NSMutableData dataWithLength: 32768];
+        NSInteger bytesRead;
+        do {
+            bytesRead = [bodyStream read: buffer.mutableBytes maxLength: buffer.length];
+            if (bytesRead > 0) {
+                [blob appendData: [NSData dataWithBytesNoCopy: buffer.mutableBytes
+                                                       length: bytesRead freeWhenDone: NO]];
+            }
+        } while (bytesRead > 0);
+        if (bytesRead < 0)
+            return kTDStatusBadAttachment;
+
+    } else {
+        NSData* body = _request.HTTPBody;
+        if (body)
+            [blob appendData: body];
+    }
+    [blob finish];
+
+    return [self updateAttachment: attachment docID: docID body: blob];
 }
 
 
 - (TDStatus) do_DELETE: (TD_Database*)db docID: (NSString*)docID attachment: (NSString*)attachment {
-    return [self updateAttachment: attachment
+    return [self updateAttachment: attachment docID: docID body: nil];
-                            docID: docID
-                             body: nil];
 }
 
 

diff --git a/Source/TD_Database+Attachments.h b/Source/TD_Database+Attachments.h
@@ -68,7 +68,7 @@ typedef enum {
 /** Updates or deletes an attachment, creating a new document revision in the process.
     Used by the PUT / DELETE methods called on attachment URLs. */
 - (TD_Revision*) updateAttachment: (NSString*)filename
-                            body: (NSData*)body
+                            body: (TDBlobStoreWriter*)body
                             type: (NSString*)contentType
                         encoding: (TDAttachmentEncoding)encoding
                          ofDocID: (NSString*)docID

diff --git a/Source/TD_Database+Attachments.m b/Source/TD_Database+Attachments.m
@@ -69,6 +69,13 @@ - (void) rememberPendingKey: (TDBlobKey)key forDigest: (NSString*)digest {
 }
 
 
+- (void) rememberAttachmentWriter: (TDBlobStoreWriter*)writer forDigest:(NSString*)digest {
+    if (!_pendingAttachmentsByDigest)
+        _pendingAttachmentsByDigest = [[NSMutableDictionary alloc] init];
+    _pendingAttachmentsByDigest[digest] = writer;
+}
+
+
 // This is ONLY FOR TESTS (see TDMultipartDownloader.m)
 #if DEBUG
 - (id) attachmentWriterForAttachment: (NSDictionary*)attachment {
@@ -561,7 +568,7 @@ - (TDMultipartWriter*) multipartWriterForRevision: (TD_Revision*)rev
 
 
 - (TD_Revision*) updateAttachment: (NSString*)filename
-                            body: (NSData*)body
+                            body: (TDBlobStoreWriter*)body
                             type: (NSString*)contentType
                         encoding: (TDAttachmentEncoding)encoding
                          ofDocID: (NSString*)docID
@@ -593,14 +600,10 @@ - (TD_Revision*) updateAttachment: (NSString*)filename
     if (!attachments)
         attachments = $mdict();
     if (body) {
-        TDBlobKey key;
+        TDBlobKey key = body.blobKey;
-        if (![self storeBlob: body creatingKey: &key]) {
-            *outStatus = kTDStatusAttachmentError;
-            return nil;
-        }
         NSString* digest = [@"sha1-" stringByAppendingString: [TDBase64 encode: &key
                                                                         length: sizeof(key)]];
-        [self rememberPendingKey: key forDigest: digest];
+        [self rememberAttachmentWriter: body forDigest: digest];
         NSString* encodingName = (encoding == kTDAttachmentEncodingGZIP) ? @"gzip" : nil;
         attachments[filename] = $dict({@"digest", digest},
                                       {@"length", @(body.length)},

diff --git a/Source/TD_Database_Tests.m b/Source/TD_Database_Tests.m
@@ -574,6 +574,14 @@ static void insertAttachment(TD_Database* db, NSData* blob,
 }
 
 
+static TDBlobStoreWriter* blobForData(TD_Database* db, NSData* data) {
+    TDBlobStoreWriter* blob = db.attachmentWriter;
+    [blob appendData: data];
+    [blob finish];
+    return blob;
+}
+
+
 TestCase(TD_Database_PutAttachment) {
     RequireTestCase(TD_Database_Attachments);
     // Start with a fresh database in /tmp:
@@ -607,17 +615,20 @@ static void insertAttachment(TD_Database* db, NSData* blob,
 
     // Update the attachment directly:
     NSData* attachv2 = [@"Replaced body of attach" dataUsingEncoding: NSUTF8StringEncoding];
-    [db updateAttachment: @"attach" body: attachv2 type: @"application/foo"
+    [db updateAttachment: @"attach" body: blobForData(db, attachv2)
+                    type: @"application/foo"
                 encoding: kTDAttachmentEncodingNone
                  ofDocID: rev1.docID revID: nil
                   status: &status];
     CAssertEq(status, kTDStatusConflict);
-    [db updateAttachment: @"attach" body: attachv2 type: @"application/foo"
+    [db updateAttachment: @"attach" body: blobForData(db, attachv2)
+                    type: @"application/foo"
                 encoding: kTDAttachmentEncodingNone
                  ofDocID: rev1.docID revID: @"1-bogus"
                   status: &status];
     CAssertEq(status, kTDStatusConflict);
-    TD_Revision* rev2 = [db updateAttachment: @"attach" body: attachv2 type: @"application/foo"
+    TD_Revision* rev2 = [db updateAttachment: @"attach" body: blobForData(db, attachv2)
+                                        type: @"application/foo"
                                    encoding: kTDAttachmentEncodingNone
                                     ofDocID: rev1.docID revID: rev1.revID
                                      status: &status];
@@ -633,6 +644,10 @@ static void insertAttachment(TD_Database* db, NSData* blob,
                                                          {@"length", @(23)},
                                                          {@"stub", $true},
                                                          {@"revpos", @2})}));
+
+    NSData* gotAttach = [db getAttachmentForSequence: gotRev2.sequence named: @"attach"
+                                                type: NULL encoding: NULL status: &status];
+    CAssertEqual(gotAttach, attachv2);
 
     // Delete the attachment:
     [db updateAttachment: @"nosuchattach" body: nil type: nil