From 7a28916ace81de81af26fb6d6d580d0fe1c02b64 Mon Sep 17 00:00:00 2001
From: Uday Alla <ualla@g.ucla.edu>
Date: Fri, 29 Oct 2021 14:50:44 -0400
Subject: [PATCH] fix: remove infinite values from message data

---
 UnitTests/MPDataModelTests.m               | 39 ++++++++++++++++++++++
 mParticle-Apple-SDK/Data Model/MPMessage.h |  3 ++
 mParticle-Apple-SDK/Data Model/MPMessage.m | 34 ++++++++++++++-----
 3 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/UnitTests/MPDataModelTests.m b/UnitTests/MPDataModelTests.m
index 5abbb2ad..b340d57a 100644
--- a/UnitTests/MPDataModelTests.m
+++ b/UnitTests/MPDataModelTests.m
@@ -102,6 +102,45 @@ - (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);
+    XCTAssertEqualObjects([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");
+}
+
 - (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
index ec7f0cf4..c1e32c56 100644
--- 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
index 32d1049b..827129b1 100644
--- 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