Improving file type detection for attachments

Changing the extension source of truth order to filename query param, MIME type, url.pathExtension, and then a new fallback of doing a reverse order query param check for any valid extension.

Updated the unit tests for this area

NSURL+OneSignal has the logic for query param file extension extraction

Author: emawby

iOS_SDK/OneSignalSDK/Source/NSURL+OneSignal.h

@@ -31,4 +31,6 @@
 
 - (NSString *)valueFromQueryParameter:(NSString *)parameter;
 
+- (NSString*)supportedFileExtension;
+
 @end

iOS_SDK/OneSignalSDK/Source/NSURL+OneSignal.m

@@ -26,6 +26,7 @@
  */
 
 #import "NSURL+OneSignal.h"
+#import "OneSignalCommonDefines.h"
 
 @implementation NSURL (OneSignal)
 - (NSString *)valueFromQueryParameter:(NSString *)parameter {
@@ -37,4 +38,28 @@ - (NSString *)valueFromQueryParameter:(NSString *)parameter {
 
     return nil;
 }
+
+- (NSString*)supportedFileExtension {
+    NSURLComponents *components = [NSURLComponents componentsWithURL:self resolvingAgainstBaseURL:false];
+    
+    for(int i = (int)components.queryItems.count-1; i > -1; i--) {
+        NSURLQueryItem *item = components.queryItems[i];
+        NSString *value = item.value;
+        NSString *extension = [self findExtensionInParam:value];
+        if (extension)
+            return extension;
+    }
+    return nil;
+}
+
+- (NSString*)findExtensionInParam:(NSString *)parameter {
+    NSArray *paramComponents = [parameter componentsSeparatedByString:@"."];
+    for (int i = (int)paramComponents.count-1; i > -1; i--) {
+        NSString *component = paramComponents[i];
+        if ([ONESIGNAL_SUPPORTED_ATTACHMENT_TYPES containsObject:component])
+            return component;
+    }
+    return nil;
+}
+
 @end

iOS_SDK/OneSignalSDK/Source/OneSignalHelper.m

@@ -81,6 +81,11 @@ -(void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataT
 
 -(void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask didReceiveResponse:(NSURLResponse *)aResponse completionHandler:(void (^)(NSURLSessionResponseDisposition))completionHandler {
     response = aResponse;
+    long long expectedLength = response.expectedContentLength;
+    if (expectedLength > 50000000) { //Enforcing 50 mb limit on media before downloading
+        completionHandler(NSURLSessionResponseCancel);
+        return;
+    }
     completionHandler(NSURLSessionResponseAllow);
 }
 
@@ -785,21 +790,16 @@ + (void)addNotificationRequest:(OSNotificationPayload*)payload
 + (NSString*)downloadMediaAndSaveInBundle:(NSString*)urlString {
 
     let url = [NSURL URLWithString:urlString];
-    
-    NSString* extension = url.pathExtension;
-    
-    if ([extension isEqualToString:@""])
-        extension = nil;
-    
-    // Unrecognized extention
-    if (extension != nil && ![ONESIGNAL_SUPPORTED_ATTACHMENT_TYPES containsObject:extension])
-        return nil;
-    
+     
+    //Try to get extension from the filname parameter
+    NSString* extension = [[[NSURL URLWithString:urlString] valueFromQueryParameter:@"filename"]
+                            supportedFileExtension];
+     
+    //Download the file
     var name = [self randomStringWithLength:10];
-    
+        
     if (extension)
         name = [name stringByAppendingString:[NSString stringWithFormat:@".%@", extension]];
-    
     NSArray* paths = NSSearchPathForDirectoriesInDomains(NSCachesDirectory, NSUserDomainMask, YES);
     NSString* filePath = [paths[0] stringByAppendingPathComponent:name];
 
@@ -813,18 +813,24 @@ + (NSString*)downloadMediaAndSaveInBundle:(NSString*)urlString {
             [OneSignal onesignal_Log:ONE_S_LL_ERROR message:[NSString stringWithFormat:@"Encountered an error while attempting to download file with URL: %@", error]];
             return nil;
         }
-        
         if (!extension) {
             NSString *newExtension;
-            
+            //Use the MIME type for the extension if one wasn't provided in the filename parameter
             if (mimeType != nil && ![mimeType isEqualToString:@""]) {
                 newExtension = mimeType.fileExtensionForMimeType;
-            } else {
-                newExtension = [[[NSURL URLWithString:urlString] valueFromQueryParameter:@"filename"] supportedFileExtension];
             }
 
-            if (!newExtension || ![ONESIGNAL_SUPPORTED_ATTACHMENT_TYPES containsObject:newExtension])
-                return nil;
+            //Try using url.pathExtension
+            if (!newExtension || ![ONESIGNAL_SUPPORTED_ATTACHMENT_TYPES containsObject:newExtension]) {
+                newExtension =  url.pathExtension;
+            }
+            
+            //Try getting an extension from the query
+            if (!newExtension || ![ONESIGNAL_SUPPORTED_ATTACHMENT_TYPES containsObject:newExtension]) {
+                newExtension = url.supportedFileExtension;
+                if (!newExtension || [newExtension isEqualToString:@""])
+                     return nil;
+            }
 
             name = [NSString stringWithFormat:@"%@.%@", name, newExtension];
 
@@ -834,12 +840,12 @@ + (NSString*)downloadMediaAndSaveInBundle:(NSString*)urlString {
         }
 
         if (error) {
-            [OneSignal onesignal_Log:ONE_S_LL_ERROR message:[NSString stringWithFormat:@"Encountered an error while attempting to download file with URL: %@", error]];
+            [OneSignal onesignal_Log:ONE_S_LL_ERROR message:[NSString stringWithFormat:@"Encountered an error    while attempting to download file with URL: %@", error]];
             return nil;
         }
-        
+         
         let standardUserDefaults = OneSignalUserDefaults.initStandard;
-        
+         
         NSArray* cachedFiles = [standardUserDefaults getSavedObjectForKey:OSUD_TEMP_CACHED_NOTIFICATION_MEDIA defaultValue:nil];
         NSMutableArray* appendedCache;
         if (cachedFiles) {
@@ -848,15 +854,14 @@ + (NSString*)downloadMediaAndSaveInBundle:(NSString*)urlString {
         }
         else
             appendedCache = [[NSMutableArray alloc] initWithObjects:name, nil];
-        
+         
         [standardUserDefaults saveObjectForKey:OSUD_TEMP_CACHED_NOTIFICATION_MEDIA withValue:appendedCache];
         return name;
     } @catch (NSException *exception) {
         [OneSignal onesignal_Log:ONE_S_LL_ERROR message:[NSString stringWithFormat:@"OneSignal encountered an exception while downloading file (%@), exception: %@", url, exception.description]];
 
         return nil;
     }
-
 }
 
 // TODO: Add back after testing

iOS_SDK/OneSignalSDK/UnitTests/Shadows/NSURLSessionOverrider.m

@@ -50,8 +50,8 @@ + (NSString *)overrideDownloadItemAtURL:(NSURL*)url toFile:(NSString*)localPath
 
     if ([url.absoluteString isEqualToString:@"http://domain.com/file"])
         return @"image/png";
-    else if ([url.absoluteString isEqualToString:@"http://domain.com/secondFile"])
-        return nil;
+    else if ([url.path isEqualToString:@"/secondFile"])
+        return @"image/heic";
     else
         return @"image/jpg";
 }

iOS_SDK/OneSignalSDK/UnitTests/UnitTests.m

@@ -2136,16 +2136,7 @@ - (id)exampleNotificationJSONWithMediaURL:(NSString *)urlString {
                      }};
 }
 
-- (void)testExtractFileExtensionFromMimeType {
-    //test to make sure the MIME type parsing works correctly
-    //NSURLSessionOverrider returns image/png for this URL
-    id pngFormat = [self exampleNotificationJSONWithMediaURL:@"http://domain.com/file"];
-    
-    let downloadedPngFilename = [self deliverNotificationWithJSON:pngFormat].URL.lastPathComponent;
-    XCTAssertTrue([downloadedPngFilename.supportedFileExtension isEqualToString:@"png"]);
-}
-
-- (void)testExtractFileExtensionFromQueryParameter {
+- (void)testExtractFileExtensionFromFileNameQueryParameter {
     // we allow developers to add ?filename=test.jpg (for example) to attachment URL's in cases where there is no extension & no mime type
     // tests to make sure the SDK correctly extracts the file extension from the `filename` URL query parameter
     // NSURLSessionOverrider returns nil for this URL
@@ -2155,14 +2146,32 @@ - (void)testExtractFileExtensionFromQueryParameter {
     XCTAssertTrue([downloadedJpgFilename.supportedFileExtension isEqualToString:@"jpg"]);
 }
 
-- (void)testFileExtensionPrioritizesURLFileExtension {
-    //tests to make sure that the URL's file extension is prioritized above the MIME type and URL query param
-    //this attachment URL will have a file extension, a MIME type, and a filename query parameter. It should prioritize the URL file extension (gif)
+- (void)testExtractFileExtensionFromMimeType {
+    //test to make sure the MIME type parsing works correctly
+    //NSURLSessionOverrider returns image/png for this URL
+    id pngFormat = [self exampleNotificationJSONWithMediaURL:@"http://domain.com/file"];
+    
+    let downloadedPngFilename = [self deliverNotificationWithJSON:pngFormat].URL.lastPathComponent;
+    XCTAssertTrue([downloadedPngFilename.supportedFileExtension isEqualToString:@"png"]);
+}
+
+- (void)testFileExtensionPrioritizesFileNameParameter {
+    //tests to make sure that the filename query parameter is prioritized above the MIME type and URL extension
+    //this attachment URL will have a file extension, a MIME type, and a filename query parameter. It should prioritize the filename query parameter (png)
     //NSURLSessionOverrider returns image/png for this URL
     id gifFormat = [self exampleNotificationJSONWithMediaURL:@"http://domain.com/file.gif?filename=test.png"];
 
     let downloadedGifFilename = [self deliverNotificationWithJSON:gifFormat].URL.lastPathComponent;
-    XCTAssertTrue([downloadedGifFilename.supportedFileExtension isEqualToString:@"gif"]);
+    XCTAssertTrue([downloadedGifFilename.supportedFileExtension isEqualToString:@"png"]);
+}
+
+- (void)testExtractFileExtensionFromAnyParamter {
+    //test to make sure the fallback of parsing all parameters for a file type works correctly
+    //NSURLSessionOverrider returns an unallowed extension (heic) for this URL to test the fallback
+    id pngFormat = [self exampleNotificationJSONWithMediaURL:@"http://domain.com/secondFile?file=test.png&media=image&type=.fakeextension"];
+    
+    let downloadedPngFilename = [self deliverNotificationWithJSON:pngFormat].URL.lastPathComponent;
+    XCTAssertTrue([downloadedPngFilename.supportedFileExtension isEqualToString:@"png"]);
 }
 
 /*