Created an way to attach an image URL to an image for a replay

Agent.xcodeproj/project.pbxproj

@@ -1576,6 +1576,8 @@
 		F8A455492AFBE31E0057B1E0 /* NRMAURLSessionHeaderTrackingTestsOldEventSystem.m in Sources */ = {isa = PBXBuildFile; fileRef = F8A455472AFBE31E0057B1E0 /* NRMAURLSessionHeaderTrackingTestsOldEventSystem.m */; };
 		F8AC3E932938FD6C002B4AA8 /* NRMAFakeDataHelper.m in Sources */ = {isa = PBXBuildFile; fileRef = F8AC3E922938FD6C002B4AA8 /* NRMAFakeDataHelper.m */; };
 		F8AC3E942938FD6C002B4AA8 /* NRMAFakeDataHelper.m in Sources */ = {isa = PBXBuildFile; fileRef = F8AC3E922938FD6C002B4AA8 /* NRMAFakeDataHelper.m */; };
+		F8CCE8752EBBA6560042C230 /* NRMAUIImageOverride.m in Sources */ = {isa = PBXBuildFile; fileRef = F8CCE8732EBBA6560042C230 /* NRMAUIImageOverride.m */; };
+		F8CCE8762EBBA6560042C230 /* NRMAUIImageOverride.h in Headers */ = {isa = PBXBuildFile; fileRef = F8CCE8722EBBA6560042C230 /* NRMAUIImageOverride.h */; };
 		F8E17C542DB681820098C3CB /* NRLogger.swift in Sources */ = {isa = PBXBuildFile; fileRef = F8E17C532DB681820098C3CB /* NRLogger.swift */; };
 		F8E17C552DB681820098C3CB /* NRLogger.swift in Sources */ = {isa = PBXBuildFile; fileRef = F8E17C532DB681820098C3CB /* NRLogger.swift */; };
 		F8E17C562DB681820098C3CB /* NRLogger.swift in Sources */ = {isa = PBXBuildFile; fileRef = F8E17C532DB681820098C3CB /* NRLogger.swift */; };
@@ -2613,6 +2615,8 @@
 		F8A455472AFBE31E0057B1E0 /* NRMAURLSessionHeaderTrackingTestsOldEventSystem.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = NRMAURLSessionHeaderTrackingTestsOldEventSystem.m; sourceTree = "<group>"; };
 		F8AC3E922938FD6C002B4AA8 /* NRMAFakeDataHelper.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = NRMAFakeDataHelper.m; sourceTree = "<group>"; };
 		F8AC3EA72938FDDB002B4AA8 /* NRMAFakeDataHelper.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = NRMAFakeDataHelper.h; sourceTree = "<group>"; };
+		F8CCE8722EBBA6560042C230 /* NRMAUIImageOverride.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = NRMAUIImageOverride.h; sourceTree = "<group>"; };
+		F8CCE8732EBBA6560042C230 /* NRMAUIImageOverride.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = NRMAUIImageOverride.m; sourceTree = "<group>"; };
 		F8E17C532DB681820098C3CB /* NRLogger.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = NRLogger.swift; sourceTree = "<group>"; };
 		F8E202DD2B07BA61008E0B7B /* NRMAOfflineStorage.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = NRMAOfflineStorage.m; sourceTree = "<group>"; };
 		F8E202F32B07BA6E008E0B7B /* NRMAOfflineStorage.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = NRMAOfflineStorage.h; sourceTree = "<group>"; };
@@ -3370,6 +3374,7 @@
 				02FF486C24DC618F00115469 /* NRMAThread.h */,
 				02FF486E24DC618F00115469 /* NRMAThread.m */,
 				02FF486B24DC618F00115469 /* NRMAThreadTransition.h */,
+				F8CCE8742EBBA6560042C230 /* UIImage */,
 				02FF486A24DC618F00115469 /* NRMAThreadTransition.m */,
 			);
 			path = Instrumentation;
@@ -4152,6 +4157,15 @@
 			path = AttributeValidator;
 			sourceTree = "<group>";
 		};
+		F8CCE8742EBBA6560042C230 /* UIImage */ = {
+			isa = PBXGroup;
+			children = (
+				F8CCE8722EBBA6560042C230 /* NRMAUIImageOverride.h */,
+				F8CCE8732EBBA6560042C230 /* NRMAUIImageOverride.m */,
+			);
+			path = UIImage;
+			sourceTree = "<group>";
+		};
 		F8FBFA3C2A71A32400CDC8C5 /* Events */ = {
 			isa = PBXGroup;
 			children = (
@@ -4348,6 +4362,7 @@
 				02FF493924DC625A00115469 /* NRMAHexUploader.h in Headers */,
 				02FF48EA24DC622700115469 /* NRMATimestampContainer.h in Headers */,
 				02FF4A6D24DC64E600115469 /* NRMAActivityTraceMeasurementProducer.h in Headers */,
+				F8CCE8762EBBA6560042C230 /* NRMAUIImageOverride.h in Headers */,
 				02FF49F024DC645B00115469 /* NRMAAppUpgradeMetricGenerator.h in Headers */,
 				02FF484124DC614200115469 /* NRMATraceMachineAgentUserInterface.h in Headers */,
 				2B496B1C2C530A7800A0459E /* reflection_generated.h in Headers */,
@@ -5681,6 +5696,7 @@
 				2BAE5C782E85FA62001D2B88 /* SwiftUIDrawingThingy.swift in Sources */,
 				02FF49CB24DC62B800115469 /* NRMAHarvestableMetric.m in Sources */,
 				02FF489A24DC61D200115469 /* NRMAURLSessionTaskDelegate.m in Sources */,
+				F8CCE8752EBBA6560042C230 /* NRMAUIImageOverride.m in Sources */,
 				02FF482224DB5B8100115469 /* NRMATableViewIntrumentation.m in Sources */,
 				02FF482124DB5B8100115469 /* NRMAApplicationInstrumentation.m in Sources */,
 				2B81C3442E996304002F5593 /* MaskedContainerView.swift in Sources */,

Agent/APrivateHeader.h

@@ -22,5 +22,6 @@
 #import "NRMABool.h"
 #import "Constants.h"
 #import "NRMAExceptionMetaDataStore.h"
+#import "NRMAUIImageOverride.h"
 
 #endif /* APrivateHeader_h */

Agent/Instrumentation/NSURLSession/NRMAURLSessionOverride.m

@@ -20,6 +20,8 @@
 #import "NRMAAssociate.h"
 #import "NRMAURLSessionTaskSearch.h"
 #import "NRMAFlags.h"
+#import "NRMAUIImageOverride.h"
+#import "NewRelicAgentInternal.h"
 
 #define NRMASwizzledMethodPrefix @"_NRMAOverride__"
 
@@ -295,7 +297,12 @@ + (void)swizzleURLSessionTask
         }
 
         // NRLOG_AGENT_VERBOSE(@"NRMA__recordTask called from NRMAOverride__dataTaskWithRequest_completionHandler");
-
+#if TARGET_OS_IOS
+        if ([[NewRelicAgentInternal sharedInstance] isSessionReplayEnabled] && [[NewRelicAgentInternal sharedInstance] isSessionReplaySampled]) {
+            [NRMAUIImageOverride registerURL:response.URL forData:data];
+        }
+#endif
+
         NRMA__recordTask(task,data,response,error);
 
         completionHandler(data,response,error);

Agent/Instrumentation/UIImage/NRMAUIImageOverride.h

@@ -0,0 +1,18 @@
+//
+//  NRMAUIImageOverride.h
+//  Agent
+//
+//  Created by Mike Bruin on 1/5/25.
+//  Copyright © 2025 New Relic. All rights reserved.
+//
+
+#import <Foundation/Foundation.h>
+#import <UIKit/UIKit.h>
+
+@interface NRMAUIImageOverride : NSObject
+
++ (void)beginInstrumentation;
++ (void)deinstrument;
++ (void)registerURL:(NSURL*)url forData:(NSData*)data;
+
+@end

Agent/Instrumentation/UIImage/NRMAUIImageOverride.m

@@ -0,0 +1,200 @@
+//
+//  NRMAUIImageOverride.m
+//  Agent
+//
+//  Created by Mike Bruin on 1/5/25.
+//  Copyright © 2025 New Relic. All rights reserved.
+//
+
+#import "NRMAUIImageOverride.h"
+#import "NRMAMethodSwizzling.h"
+#import "NRLogger.h"
+#import <objc/runtime.h>
+#import <CommonCrypto/CommonDigest.h>
+#import <NewRelic/NewRelic-Swift.h>
+
+static IMP NRMAOriginal__initWithData;
+static IMP NRMAOriginal__initWithData_scale;
+
+// Global registry to map data hashes to URLs
+static NSMapTable *dataHashToURLMap;
+static NSLock *registryLock;
+
+// Flag to prevent double swizzling
+static BOOL isSwizzled = NO;
+
+// Forward declarations
+UIImage* NRMAOverride__initWithData(UIImage* self, SEL _cmd, NSData* data);
+UIImage* NRMAOverride__initWithData_scale(UIImage* self, SEL _cmd, NSData* data, CGFloat scale);
+NSString* NRMA_HashForData(NSData* data);
+
+@implementation NRMAUIImageOverride
+
++ (void)beginInstrumentation {
+    // Prevent double swizzling
+    if (isSwizzled) {
+        NRLOG_AGENT_DEBUG(@"NRMAUIImageOverride - Already instrumented, skipping swizzling");
+        return;
+    }
+
+    Class clazz = [UIImage class];
+
+    // Initialize the global registry and lock
+    if (dataHashToURLMap == nil) {
+        dataHashToURLMap = [NSMapTable strongToStrongObjectsMapTable];
+        registryLock = [[NSLock alloc] init];
+    }
+
+    if (clazz) {
+        // Swizzle initWithData:
+        NRMAOriginal__initWithData = NRMASwapImplementations(clazz,
+            @selector(initWithData:),
+            (IMP)NRMAOverride__initWithData);
+
+        // Swizzle initWithData:scale:
+        NRMAOriginal__initWithData_scale = NRMASwapImplementations(clazz,
+            @selector(initWithData:scale:),
+            (IMP)NRMAOverride__initWithData_scale);
+
+        // Mark as swizzled
+        isSwizzled = YES;
+        NRLOG_AGENT_DEBUG(@"NRMAUIImageOverride - Instrumentation completed successfully");
+    }
+}
+
+// Public method to register a URL for NSData
++ (void)registerURL:(NSURL*)url forData:(NSData*)data {
+    if (url == nil || data == nil) return;
+
+    NSString *hash = NRMA_HashForData(data);
+    [registryLock lock];
+
+    [dataHashToURLMap setObject:url forKey:hash];
+
+    [registryLock unlock];
+
+    NRLOG_AGENT_DEBUG(@"NRMAUIImageOverride map size: %lu", (unsigned long)dataHashToURLMap.count);
+}
+
++ (void)deinstrument {
+    Class clazz = [UIImage class];
+
+    if (clazz) {
+        if (NRMAOriginal__initWithData != nil) {
+            NRMASwapImplementations(clazz, @selector(initWithData:), (IMP)NRMAOriginal__initWithData);
+            NRMAOriginal__initWithData = nil;
+        }
+
+        if (NRMAOriginal__initWithData_scale != nil) {
+            NRMASwapImplementations(clazz, @selector(initWithData:scale:), (IMP)NRMAOriginal__initWithData_scale);
+            NRMAOriginal__initWithData_scale = nil;
+        }
+    }
+
+    // Clean up registry
+    [registryLock lock];
+    [dataHashToURLMap removeAllObjects];
+    [registryLock unlock];
+
+    // Reset the swizzled flag
+    isSwizzled = NO;
+}
+
+@end
+
+// Swizzled implementation for initWithData:
+UIImage* NRMAOverride__initWithData(UIImage* self, SEL _cmd, NSData* data) {
+    if (NRMAOriginal__initWithData == nil || data == nil) {
+        return nil;
+    }
+
+    NSString *hash = NRMA_HashForData(data);
+    [registryLock lock];
+    NSURL* url = [dataHashToURLMap objectForKey:hash];
+
+    // Remove the entry after retrieving it (automatic cleanup after use)
+    if (url != nil) {
+        [dataHashToURLMap removeObjectForKey:hash];
+    }
+
+    [registryLock unlock];
+
+    // Call original implementation
+    UIImage* image = ((UIImage*(*)(id, SEL, NSData*))NRMAOriginal__initWithData)(self, _cmd, data);
+
+    if (image != nil && url != nil) {
+        // Attach the URL to the newly created UIImage
+        image.NRSessionReplayImageURL = url;
+        NRLOG_AGENT_DEBUG(@"NRMAOverride__initWithData - Successfully attached URL to image: %@", url);
+    }
+
+    return image;
+}
+
+// Swizzled implementation for initWithData:scale:
+UIImage* NRMAOverride__initWithData_scale(UIImage* self, SEL _cmd, NSData* data, CGFloat scale) {
+    if (NRMAOriginal__initWithData_scale == nil || data == nil) {
+        return nil;
+    }
+
+    NSString *hash = NRMA_HashForData(data);
+    [registryLock lock];
+    NSURL* url = [dataHashToURLMap objectForKey:hash];
+
+    // Remove the entry after retrieving it (automatic cleanup after use)
+    if (url != nil) {
+        [dataHashToURLMap removeObjectForKey:hash];
+    }
+
+    [registryLock unlock];
+
+    // Call original implementation
+    UIImage* image = ((UIImage*(*)(id, SEL, NSData*, CGFloat))NRMAOriginal__initWithData_scale)(self, _cmd, data, scale);
+
+    if (image != nil && url != nil) {
+        // Attach the URL to the newly created UIImage
+        image.NRSessionReplayImageURL = url;
+        NRLOG_AGENT_DEBUG(@"NRMAOverride__initWithData:scale: - Successfully attached URL to image: %@", url);
+    }
+
+    return image;
+}
+
+// Helper function to generate a hash for NSData
+NSString* NRMA_HashForData(NSData* data) {
+    if (data == nil || data.length == 0) {
+        return nil;
+    }
+
+    // For performance, only hash the first 1KB and last 1KB along with the length
+    // This is sufficient to uniquely identify image data in most cases
+    NSUInteger length = data.length;
+    NSUInteger hashLength = MIN(1024, length);
+
+    unsigned char hash[CC_SHA256_DIGEST_LENGTH];
+    CC_SHA256_CTX ctx;
+    CC_SHA256_Init(&ctx);
+
+    // Hash the length
+    CC_SHA256_Update(&ctx, &length, sizeof(length));
+
+    // Hash first chunk
+    CC_SHA256_Update(&ctx, data.bytes, (CC_LONG)hashLength);
+
+    // If data is large enough, also hash last chunk
+    if (length > 2048) {
+        const void *lastChunk = data.bytes + (length - hashLength);
+        CC_SHA256_Update(&ctx, lastChunk, (CC_LONG)hashLength);
+    }
+
+    CC_SHA256_Final(hash, &ctx);
+
+    // Convert to hex string
+    NSMutableString *hashString = [NSMutableString stringWithCapacity:CC_SHA256_DIGEST_LENGTH * 2];
+    for (int i = 0; i < CC_SHA256_DIGEST_LENGTH; i++) {
+        [hashString appendFormat:@"%02x", hash[i]];
+    }
+
+    return hashString;
+}
+

Agent/SessionReplay/NRMASessionReplay.swift

@@ -59,6 +59,8 @@ public class NRMASessionReplay: NSObject {
                 }
                 self.sessionReplayTouchCapture = SessionReplayTouchCapture(window: window)
                 swizzleSendEvent()
+                // Instrument UIImage to preserve URL from NSData
+                NRMAUIImageOverride.beginInstrumentation()
             }
         }
     }