fix: Prevent Infinite Values in Message Data (#39)

mParticle · Oct 4, 2021 · d19d785 · d19d785
1 parent 8a3437f
commit d19d785
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 11 deletions.
diff --git a/UnitTests/MPDataModelTests.m b/UnitTests/MPDataModelTests.m
@@ -102,6 +102,42 @@ - (void)testMessageInstance {
     XCTAssertNotNil(dictionaryRepresentation, @"Should not have been nil.");
 }
 
+- (void)testMessageInstanceWithInfinite {
+    MPSession *session = [[MPSession alloc] initWithStartTime:[[NSDate date] timeIntervalSince1970] userId:[MPPersistenceController mpId]];
+
+    double four = 4.0;
+    double zed = 0.0;
+    MPMessageBuilder *messageBuilder = [MPMessageBuilder newBuilderWithMessageType:MPMessageTypeEvent
+                                                                           session:session
+                                                                       messageInfo:@{@"MessageKey1":@(four/zed)}];
+    XCTAssertNotNil(messageBuilder, @"Should not have been nil.");
+
+    MPMessage *message = [messageBuilder build];
+    XCTAssertNotNil(message, @"Should not have been nil.");
+
+    NSString *description = [message description];
+    XCTAssertNotNil(description, @"Should not have been nil.");
+
+    MPMessage *messageCopy = [message copy];
+    XCTAssertNotNil(messageCopy, @"Should not have been nil.");
+    XCTAssertEqualObjects(message, messageCopy, @"Should have been equal.");
+    messageCopy.timestamp = [[NSDate date] timeIntervalSince1970];
+    XCTAssertNotEqualObjects(message, messageCopy, @"Should not have been equal.");
+
+    messageCopy = (MPMessage *)[NSNull null];
+    XCTAssertNotEqualObjects(message, messageCopy, @"Should not have been equal.");
+    XCTAssertNotEqualObjects(messageCopy, message, @"Should not have been equal.");
+    messageCopy = (MPMessage *)@"This is not a valid message object.";
+    XCTAssertNotEqualObjects(message, messageCopy, @"Should not have been equal.");
+    XCTAssertNotEqualObjects(messageCopy, message, @"Should not have been equal.");
+    messageCopy = nil;
+    XCTAssertNotEqualObjects(message, messageCopy, @"Should not have been equal.");
+    XCTAssertNotEqualObjects(messageCopy, message, @"Should not have been equal.");
+
+    NSDictionary *dictionaryRepresentation = [message dictionaryRepresentation];
+    XCTAssertNotNil(dictionaryRepresentation, @"Should not have been nil.");
+}
+
 - (void)testUploadInstance {
     MPSession *session = [[MPSession alloc] initWithStartTime:[[NSDate date] timeIntervalSince1970] userId:[MPPersistenceController mpId]];
 

diff --git a/mParticle-Apple-SDK/Data Model/MPMessage.m b/mParticle-Apple-SDK/Data Model/MPMessage.m
@@ -39,11 +39,23 @@ - (instancetype)initWithSession:(MPSession *)session messageType:(NSString *)mes
         sessionId = @(session.sessionId);
     }
 
+    NSMutableDictionary *messageDictionary = messageInfo.mutableCopy;
+    for (NSString *key in messageInfo) {
+        if ([messageInfo[key] isKindOfClass:[NSNumber class]]) {
+            NSNumber *value = (NSNumber *)messageInfo[key];
+            if(value.doubleValue == INFINITY || value.doubleValue == -INFINITY || isnan(value.doubleValue)) {
+                MPILogVerbose(@"Invalid Message Data for key: %@", key);
+                MPILogVerbose(@"Value should not be infinite. Removing value from message data");
+                messageDictionary[key] = nil;
+            }
+        }
+    }
+
     return [self initWithSessionId:sessionId
                          messageId:0
                               UUID:uuid
                        messageType:messageType
-                       messageData:[NSJSONSerialization dataWithJSONObject:messageInfo options:0 error:nil]
+                       messageData:[NSJSONSerialization dataWithJSONObject:messageDictionary options:0 error:nil]
                          timestamp:timestamp
                       uploadStatus:uploadStatus
                             userId:userId

diff --git a/mParticle-Apple-SDK/Utils/MPMessageBuilder.mm b/mParticle-Apple-SDK/Utils/MPMessageBuilder.mm
@@ -382,16 +382,6 @@ - (MPMessage *)build {
     messageDictionary[kMPMessageIdKey] = uuid ? uuid : [[NSUUID UUID] UUIDString];
 
     NSNumber *userId = _session.userId.integerValue ? _session.userId : [MPPersistenceController mpId];
-    for (NSString *key in messageDictionary) {
-        if ([messageDictionary[key] isKindOfClass:[NSNumber class]]) {
-            NSNumber *value = (NSNumber *)messageDictionary[key];
-            if(value.doubleValue == INFINITY) {
-                MPILogVerbose(@"Invalid Message Data for key: %@", key);
-                MPILogVerbose(@"Value should not be infinite. Removing value from message data");
-                messageDictionary[key] = nil;
-            }
-        }
-    }
 
     message = [[MPMessage alloc] initWithSession:_session
                                      messageType:_messageType