Include size in attachment pointer

// FREEBIE

Signal/src/ViewControllers/DebugUI/DebugUIMessages.m

@@ -917,10 +917,12 @@ NS_ASSUME_NONNULL_BEGIN
                     break;
                 }
                 case 2: {
+                    UInt32 filesize = 64;
                     TSAttachmentPointer *pointer =
                         [[TSAttachmentPointer alloc] initWithServerId:237391539706350548
-                                                                  key:[self createRandomNSDataOfSize:64]
+                                                                  key:[self createRandomNSDataOfSize:filesize]
                                                                digest:nil
+                                                            byteCount:filesize
                                                           contentType:@"audio/mp3"
                                                                 relay:@""
                                                        sourceFilename:@"test.mp3"
@@ -947,11 +949,14 @@ NS_ASSUME_NONNULL_BEGIN
                                                     expiresInSeconds:0];
 
                     NSString *filename = @"test.mp3";
-                    TSAttachmentStream *attachmentStream =
+                    UInt32 filesize = 16;
-                        [[TSAttachmentStream alloc] initWithContentType:@"audio/mp3" sourceFilename:filename];
+
+                    TSAttachmentStream *attachmentStream = [[TSAttachmentStream alloc] initWithContentType:@"audio/mp3"
+                                                                                                 byteCount:filesize
+                                                                                            sourceFilename:filename];
 
                     NSError *error;
-                    BOOL success = [attachmentStream writeData:[self createRandomNSDataOfSize:16] error:&error];
+                    BOOL success = [attachmentStream writeData:[self createRandomNSDataOfSize:filesize] error:&error];
                     OWSAssert(success && !error);
 
                     [attachmentStream saveWithTransaction:transaction];

SignalServiceKit/src/Messages/Attachments/OWSAttachmentsProcessor.m

@@ -98,6 +98,7 @@ static const CGFloat kAttachmentDownloadProgressTheta = 0.001f;
         TSAttachmentPointer *pointer = [[TSAttachmentPointer alloc] initWithServerId:attachmentProto.id
                                                                                  key:attachmentProto.key
                                                                               digest:digest
+                                                                           byteCount:attachmentProto.size
                                                                          contentType:attachmentProto.contentType
                                                                                relay:relay
                                                                       sourceFilename:attachmentProto.fileName

SignalServiceKit/src/Messages/Attachments/TSAttachment.h

@@ -27,10 +27,12 @@ typedef NS_ENUM(NSUInteger, TSAttachmentType) {
 @property (atomic, readwrite) UInt64 serverId;
 @property (atomic, readwrite) NSData *encryptionKey;
 @property (nonatomic, readonly) NSString *contentType;
-
 @property (atomic, readwrite) BOOL isDownloaded;
 @property (nonatomic) TSAttachmentType attachmentType;
 
+// Though now required, may incorrectly be 0 on legacy attachments.
+@property (nonatomic, readonly) UInt32 byteCount;
+
 // Represents the "source" filename sent or received in the protos,
 // not the filename on disk.
 @property (nonatomic, readonly, nullable) NSString *sourceFilename;
@@ -39,12 +41,15 @@ typedef NS_ENUM(NSUInteger, TSAttachmentType) {
 // i.e. undownloaded incoming attachments.
 - (instancetype)initWithServerId:(UInt64)serverId
                    encryptionKey:(NSData *)encryptionKey
+                       byteCount:(UInt32)byteCount
                      contentType:(NSString *)contentType
                   sourceFilename:(nullable NSString *)sourceFilename;
 
 // This constructor is used for new instances of TSAttachmentStream
 // that represent new, un-uploaded outgoing attachments.
-- (instancetype)initWithContentType:(NSString *)contentType sourceFilename:(nullable NSString *)sourceFilename;
+- (instancetype)initWithContentType:(NSString *)contentType
+                          byteCount:(UInt32)byteCount
+                     sourceFilename:(nullable NSString *)sourceFilename;
 
 // This constructor is used for new instances of TSAttachmentStream
 // that represent downloaded incoming attachments.

SignalServiceKit/src/Messages/Attachments/TSAttachment.m

@@ -21,9 +21,18 @@ NSUInteger const TSAttachmentSchemaVersion = 3;
 // i.e. undownloaded incoming attachments.
 - (instancetype)initWithServerId:(UInt64)serverId
                    encryptionKey:(NSData *)encryptionKey
+                       byteCount:(UInt32)byteCount
                      contentType:(NSString *)contentType
                   sourceFilename:(nullable NSString *)sourceFilename
 {
+    OWSAssert(serverId > 0);
+    OWSAssert(encryptionKey.length > 0);
+    if (byteCount <= 0) {
+        // This will fail with legacy iOS clients which don't upload attachment size.
+        DDLogWarn(@"%@ Missing byteCount for attachment with serverId: %lld", self.tag, serverId);
+    }
+    OWSAssert(contentType.length > 0);
+
     self = [super init];
     if (!self) {
         return self;
@@ -31,26 +40,35 @@ NSUInteger const TSAttachmentSchemaVersion = 3;
 
     _serverId = serverId;
     _encryptionKey = encryptionKey;
+    _byteCount = byteCount;
     _contentType = contentType;
-    _attachmentSchemaVersion = TSAttachmentSchemaVersion;
     _sourceFilename = sourceFilename;
 
+    _attachmentSchemaVersion = TSAttachmentSchemaVersion;
+
     return self;
 }
 
 // This constructor is used for new instances of TSAttachmentStream
 // that represent new, un-uploaded outgoing attachments.
-- (instancetype)initWithContentType:(NSString *)contentType sourceFilename:(nullable NSString *)sourceFilename
+- (instancetype)initWithContentType:(NSString *)contentType
+                          byteCount:(UInt32)byteCount
+                     sourceFilename:(nullable NSString *)sourceFilename
 {
+    OWSAssert(contentType.length > 0);
+    OWSAssert(byteCount > 0);
+
     self = [super init];
     if (!self) {
         return self;
     }
 
     _contentType = contentType;
-    _attachmentSchemaVersion = TSAttachmentSchemaVersion;
+    _byteCount = byteCount;
     _sourceFilename = sourceFilename;
 
+    _attachmentSchemaVersion = TSAttachmentSchemaVersion;
+
     return self;
 }
 
@@ -58,6 +76,14 @@ NSUInteger const TSAttachmentSchemaVersion = 3;
 // that represent downloaded incoming attachments.
 - (instancetype)initWithPointer:(TSAttachment *)pointer
 {
+    OWSAssert(pointer.serverId > 0);
+    OWSAssert(pointer.encryptionKey.length > 0);
+    if (pointer.byteCount <= 0) {
+        // This will fail with legacy iOS clients which don't upload attachment size.
+        DDLogWarn(@"%@ Missing pointer.byteCount for attachment with serverId: %lld", self.tag, pointer.serverId);
+    }
+    OWSAssert(pointer.contentType.length > 0);
+
     // Once saved, this AttachmentStream will replace the AttachmentPointer in the attachments collection.
     self = [super initWithUniqueId:pointer.uniqueId];
     if (!self) {
@@ -66,8 +92,10 @@ NSUInteger const TSAttachmentSchemaVersion = 3;
 
     _serverId = pointer.serverId;
     _encryptionKey = pointer.encryptionKey;
+    _byteCount = pointer.byteCount;
     _contentType = pointer.contentType;
     _sourceFilename = pointer.sourceFilename;
+
     _attachmentSchemaVersion = TSAttachmentSchemaVersion;
 
     return self;

SignalServiceKit/src/Messages/Attachments/TSAttachmentPointer.h

@@ -22,6 +22,7 @@ typedef NS_ENUM(NSUInteger, TSAttachmentPointerState) {
 - (instancetype)initWithServerId:(UInt64)serverId
                              key:(NSData *)key
                           digest:(nullable NSData *)digest
+                       byteCount:(UInt32)byteCount
                      contentType:(NSString *)contentType
                            relay:(NSString *)relay
                   sourceFilename:(nullable NSString *)sourceFilename

SignalServiceKit/src/Messages/Attachments/TSAttachmentPointer.m

@@ -28,12 +28,17 @@ NS_ASSUME_NONNULL_BEGIN
 - (instancetype)initWithServerId:(UInt64)serverId
                              key:(NSData *)key
                           digest:(nullable NSData *)digest
+                       byteCount:(UInt32)byteCount
                      contentType:(NSString *)contentType
                            relay:(NSString *)relay
                   sourceFilename:(nullable NSString *)sourceFilename
                   attachmentType:(TSAttachmentType)attachmentType
 {
-    self = [super initWithServerId:serverId encryptionKey:key contentType:contentType sourceFilename:sourceFilename];
+    self = [super initWithServerId:serverId
+                     encryptionKey:key
+                         byteCount:byteCount
+                       contentType:contentType
+                    sourceFilename:sourceFilename];
     if (!self) {
         return self;
     }

SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.h

@@ -17,7 +17,9 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (instancetype)init NS_UNAVAILABLE;
 - (instancetype)initWithContentType:(NSString *)contentType
+                          byteCount:(UInt32)byteCount
                      sourceFilename:(nullable NSString *)sourceFilename NS_DESIGNATED_INITIALIZER;
+
 - (instancetype)initWithPointer:(TSAttachmentPointer *)pointer NS_DESIGNATED_INITIALIZER;
 - (nullable instancetype)initWithCoder:(NSCoder *)coder NS_DESIGNATED_INITIALIZER;
 

SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.m

@@ -30,9 +30,11 @@ NS_ASSUME_NONNULL_BEGIN
 
 @implementation TSAttachmentStream
 
-- (instancetype)initWithContentType:(NSString *)contentType sourceFilename:(nullable NSString *)sourceFilename
+- (instancetype)initWithContentType:(NSString *)contentType
+                          byteCount:(UInt32)byteCount
+                     sourceFilename:(nullable NSString *)sourceFilename
 {
-    self = [super initWithContentType:contentType sourceFilename:sourceFilename];
+    self = [super initWithContentType:contentType byteCount:byteCount sourceFilename:sourceFilename];
     if (!self) {
         return self;
     }

SignalServiceKit/src/Messages/Interactions/TSOutgoingMessage.m

@@ -538,6 +538,7 @@ NSString *const kTSOutgoingMessageSentRecipientAll = @"kTSOutgoingMessageSentRec
     [builder setId:attachmentStream.serverId];
     [builder setContentType:attachmentStream.contentType];
     [builder setFileName:filename];
+    [builder setSize:attachmentStream.byteCount];
     [builder setKey:attachmentStream.encryptionKey];
     [builder setDigest:attachmentStream.digest];
     [builder setFlags:(self.isVoiceMessage ? OWSSignalServiceProtosAttachmentPointerFlagsVoiceMessage : 0)];

SignalServiceKit/src/Messages/OWSMessageSender.m

@@ -547,8 +547,9 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
     OWSAssert(dataSource);
 
     dispatch_async([OWSDispatch attachmentsQueue], ^{
-        TSAttachmentStream *attachmentStream =
+        TSAttachmentStream *attachmentStream = [[TSAttachmentStream alloc] initWithContentType:contentType
-            [[TSAttachmentStream alloc] initWithContentType:contentType sourceFilename:sourceFilename];
+                                                                                     byteCount:dataSource.dataLength
+                                                                                sourceFilename:sourceFilename];
         if (message.isVoiceMessage) {
             attachmentStream.attachmentType = TSAttachmentTypeVoiceMessage;
         }