fix: remove infinite values from message data (#138)

mParticle · Nov 3, 2021 · 3ea7370 · 3ea7370
1 parent 1bfd098
commit 3ea7370
Show file tree

Hide file tree

Showing 3 changed files with 80 additions and 9 deletions.
diff --git a/UnitTests/MPDataModelTests.m b/UnitTests/MPDataModelTests.m
@@ -102,6 +102,58 @@ - (void)testMessageInstance {
     XCTAssertNotNil(dictionaryRepresentation, @"Should not have been nil.");
 }
 
+- (void)testFixInvalidKeysInDictionary {
+    double four = 4.0;
+    double zed = 0.0;
+
+    NSDictionary *messageInfo = @{@"MessageKey1": @"MessageValue1",
+                                  @"MessageKey2": @"MessageValue2"};
+    NSMutableDictionary *messageDictionary = messageInfo.mutableCopy;
+    [MPMessage fixInvalidKeysInDictionary:messageDictionary messageInfo:messageInfo];
+    XCTAssertEqualObjects(messageInfo, messageDictionary);
+
+    messageInfo = @{@"MessageKey1": @"MessageValue1",
+                    @"MessageKey2": @(four/zed)};
+    messageDictionary = messageInfo.mutableCopy;
+    [MPMessage fixInvalidKeysInDictionary:messageDictionary messageInfo:messageInfo];
+    XCTAssertEqual([messageDictionary objectForKey:@"MessageKey2"], nil);
+    XCTAssertEqual([messageDictionary objectForKey:@"MessageKey1"], @"MessageValue1");
+
+    messageInfo = @{@"MessageKey1": @"MessageValue1",
+                    @"MessageKey2": @{
+                        @"NestedKey1": @(four/zed),
+                        @"NestedKey2": @"test"}};
+    messageDictionary = messageInfo.mutableCopy;
+    [MPMessage fixInvalidKeysInDictionary:messageDictionary messageInfo:messageInfo];
+    XCTAssertEqual([messageDictionary[@"MessageKey2"] objectForKey:@"NestedKey1"], nil);
+    XCTAssertEqual([messageDictionary[@"MessageKey2"] objectForKey:@"NestedKey2"], @"test");
+
+    messageInfo = @{@"MessageKey1": @"MessageValue1",
+                    @"MessageKey2": @{
+                        @"NestedKey1": @{ @"NestedKeyA": @(four/zed),
+                                          @"NestedKeyB": @"test"},
+                        @"NestedKey2": @"test"}
+
+    };
+    messageDictionary = messageInfo.mutableCopy;
+    [MPMessage fixInvalidKeysInDictionary:messageDictionary messageInfo:messageInfo];
+    XCTAssertEqual([messageDictionary[@"MessageKey2"][@"NestedKey1"] objectForKey:@"NestedKeyA"], nil);
+    XCTAssertEqual([messageDictionary[@"MessageKey2"][@"NestedKey1"] objectForKey:@"NestedKeyB"], @"test");
+
+    messageInfo = @{@"MessageKey1": @(four/zed),
+                    @"MessageKey2": @{
+                        @"NestedKey1": @{ @"NestedKeyA": @(four/zed),
+                                          @"NestedKeyB": @"test"},
+                        @"NestedKey2": @(four/zed)}};
+    messageDictionary = messageInfo.mutableCopy;
+    [MPMessage fixInvalidKeysInDictionary:messageDictionary messageInfo:messageInfo];
+    XCTAssertEqual([messageDictionary objectForKey:@"MessageKey1"], nil);
+    XCTAssertEqual([messageDictionary[@"MessageKey2"] objectForKey:@"NestedKey2"], nil);
+    XCTAssertEqual([messageDictionary[@"MessageKey2"][@"NestedKey1"] objectForKey:@"NestedKeyA"], nil);
+    XCTAssertEqual([messageDictionary[@"MessageKey2"][@"NestedKey1"] objectForKey:@"NestedKeyB"], @"test");
+
+}
+
 - (void)testMessageInstanceWithInfinite {
     MPSession *session = [[MPSession alloc] initWithStartTime:[[NSDate date] timeIntervalSince1970] userId:[MPPersistenceController mpId]];
 

diff --git a/mParticle-Apple-SDK/Data Model/MPMessage.h b/mParticle-Apple-SDK/Data Model/MPMessage.h
@@ -43,4 +43,7 @@
 - (void)truncateMessageDataProperty:(nonnull NSString *)property
                            toLength:(NSInteger)length;
 
++ (void)fixInvalidKeysInDictionary:(nonnull NSMutableDictionary*)messageDictionary
+                       messageInfo:(nonnull NSDictionary*) messageInfo;
+
 @end
diff --git a/mParticle-Apple-SDK/Data Model/MPMessage.m b/mParticle-Apple-SDK/Data Model/MPMessage.m
@@ -32,6 +32,29 @@ - (instancetype)initWithSessionId:(NSNumber *)sessionId messageId:(int64_t)messa
     return self;
 }
 
++ (void)fixInvalidKeysInDictionary:(NSMutableDictionary*)messageDictionary messageInfo:(NSDictionary*) messageInfo{
+    for (NSString *key in messageInfo) {
+        if ([messageInfo[key] isKindOfClass:[NSDictionary class]]) {
+            if (![NSJSONSerialization isValidJSONObject: messageInfo[key]]) {
+                NSMutableDictionary *temp = [messageDictionary[key] mutableCopy];
+                [MPMessage fixInvalidKeysInDictionary:temp messageInfo:messageInfo[key]];
+                messageDictionary[key] = temp;
+            }
+        } else {
+            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;
+}
+
 - (instancetype)initWithSession:(MPSession *)session messageType:(NSString *)messageType messageInfo:(NSDictionary *)messageInfo uploadStatus:(MPUploadStatus)uploadStatus UUID:(NSString *)uuid timestamp:(NSTimeInterval)timestamp userId:(NSNumber *)userId  dataPlanId:(NSString *)dataPlanId dataPlanVersion:(NSNumber *)dataPlanVersion {
     NSNumber *sessionId = nil;
 
@@ -40,15 +63,8 @@ - (instancetype)initWithSession:(MPSession *)session messageType:(NSString *)mes
     }
 
     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;
-            }
-        }
+    if (![NSJSONSerialization isValidJSONObject: messageInfo]) {
+        [MPMessage fixInvalidKeysInDictionary:messageDictionary messageInfo:messageInfo];
     }
 
     return [self initWithSessionId:sessionId