Replace URLProtocol property setting

Example/SBTUITestTunnel_Tests/MonitorTests.swift

@@ -205,11 +205,41 @@ class MonitorTests: XCTestCase {
         XCTAssert(app.stubRequestsRemoveAll())
         XCTAssert(app.monitorRequestRemoveAll())
     }
-
+    
     func testMonitorPostRequestWithHTTPBody() {
         app.monitorRequests(matching: SBTRequestMatch(url: "httpbin.org", method: "POST"))
+
+        let smallBody = String(repeating: "a", count: 100)
+
+        _ = request.dataTaskNetwork(urlString: "https://httpbin.org/post", httpMethod: "POST", httpBody: smallBody)
+        let requests = app.monitoredRequestsFlushAll()
+        XCTAssertEqual(requests.count, 1)
+        print(requests.map(\.debugDescription))
+
+        for request in requests {
+            guard let httpBody = request.request?.httpBody else {
+                XCTFail("Missing http body")
+                continue
+            }
+
+            XCTAssertEqual(String(data: httpBody, encoding: .utf8), smallBody)
+
+            XCTAssert((request.responseString()!).contains("httpbin.org"))
+            XCTAssert(request.timestamp > 0.0)
+            XCTAssert(request.requestTime > 0.0)
+        }
+
+        XCTAssert(app.stubRequestsRemoveAll())
+        XCTAssert(app.monitorRequestRemoveAll())
+    }
+
+
+    func testMonitorPostRequestWithHTTPLargeBody() {
+        app.monitorRequests(matching: SBTRequestMatch(url: "httpbin.org", method: "POST"))
+
+        let largeBody = String(repeating: "a", count: 20000)
 
-        _ = request.dataTaskNetwork(urlString: "https://httpbin.org/post", httpMethod: "POST", httpBody: "&param5=val5&param6=val6")
+        _ = request.dataTaskNetwork(urlString: "https://httpbin.org/post", httpMethod: "POST", httpBody: largeBody)
         let requests = app.monitoredRequestsFlushAll()
         XCTAssertEqual(requests.count, 1)
         print(requests.map(\.debugDescription))
@@ -220,7 +250,7 @@ class MonitorTests: XCTestCase {
                 continue
             }
 
-            XCTAssertEqual(String(data: httpBody, encoding: .utf8), "&param5=val5&param6=val6")
+            XCTAssertEqual(String(data: httpBody, encoding: .utf8), largeBody)
 
             XCTAssert((request.responseString()!).contains("httpbin.org"))
             XCTAssert(request.timestamp > 0.0)

Sources/SBTUITestTunnelCommon/NSURLRequest+HTTPBodyFix.m

@@ -17,6 +17,7 @@
 #import "include/NSURLRequest+HTTPBodyFix.h"
 #import "include/SBTUITestTunnel.h"
 #import "include/SBTSwizzleHelpers.h"
+#import "include/SBTRequestPropertyStorage.h"
 
 @implementation NSURLRequest (HTTPBodyFix)
 
@@ -37,20 +38,20 @@ @implementation NSURLRequest (HTTPBodyFix)
 
 - (NSData *)sbt_uploadHTTPBody
 {
-    return [NSURLProtocol propertyForKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:self];
+    return [SBTRequestPropertyStorage propertyForKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:self];
 }
 
 - (BOOL)sbt_isUploadTaskRequest
 {
-    return ([NSURLProtocol propertyForKey:SBTUITunneledNSURLProtocolIsUploadTaskKey inRequest:self] != nil);
+    return ([SBTRequestPropertyStorage propertyForKey:SBTUITunneledNSURLProtocolIsUploadTaskKey inRequest:self] != nil);
 }
 
 - (void)sbt_markUploadTaskRequest
 {
     NSAssert([self isKindOfClass:[NSMutableURLRequest class]], @"Attempted to mark an immutable request as an upload");
 
     if ([self isKindOfClass:[NSMutableURLRequest class]]) {
-        [NSURLProtocol setProperty:@YES forKey:SBTUITunneledNSURLProtocolIsUploadTaskKey inRequest:(NSMutableURLRequest *)self];
+        [SBTRequestPropertyStorage setProperty:@YES forKey:SBTUITunneledNSURLProtocolIsUploadTaskKey inRequest:(NSMutableURLRequest *)self];
     }
 }
 
@@ -81,17 +82,17 @@ - (NSData *)swz_HTTPBody
 
     NSData *ret = [self swz_HTTPBody];
 
-    return ret ?: [NSURLProtocol propertyForKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:self];
+    return ret ?: [SBTRequestPropertyStorage propertyForKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:self];
 }
 
 - (id)swz_copyWithZone:(NSZone *)zone
 {
     NSURLRequest *ret = [self swz_copyWithZone:zone];
 
     if ([ret isKindOfClass:[NSMutableURLRequest class]]) {
-        NSData *body = [NSURLProtocol propertyForKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:self];
+        NSData *body = [SBTRequestPropertyStorage propertyForKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:self];
         if (body) {
-            [NSURLProtocol setProperty:body forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)ret];            
+            [SBTRequestPropertyStorage setProperty:body forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)ret];
         }
     }
 
@@ -103,9 +104,9 @@ - (id)swz_mutableCopyWithZone:(NSZone *)zone
     NSMutableURLRequest *ret = [self swz_mutableCopyWithZone:zone];
 
     if ([ret isKindOfClass:[NSMutableURLRequest class]]) {
-        NSData *body = [NSURLProtocol propertyForKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:self];
+        NSData *body = [SBTRequestPropertyStorage propertyForKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:self];
         if (body) {
-            [NSURLProtocol setProperty:body forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)ret];
+            [SBTRequestPropertyStorage setProperty:body forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)ret];
         }
     }
 

Sources/SBTUITestTunnelCommon/SBTRequestPropertyStorage.m

@@ -0,0 +1,38 @@
+//
+//  SBTRequestProperty.m
+//  SBTUITestTunnelCommon
+//
+//  Created by tomas on 20/02/24.
+//
+
+#import "include/SBTRequestPropertyStorage.h"
+
+@implementation SBTRequestPropertyStorage
+
+static NSMutableDictionary *storage;
+
++ (void)initialize 
+{
+    if (self == [SBTRequestPropertyStorage class]) {
+        storage = [[NSMutableDictionary alloc] init];
+    }
+}
+
++ (void)setProperty:(id)property forKey:(nonnull NSString *)key inRequest:(nonnull NSMutableURLRequest *)request
+{
+    if ([property isKindOfClass:[NSData class]] && ((NSData *)property).length > 16834) {
+        NSString *uuid = [[NSUUID UUID] UUIDString];
+        [storage setObject:property forKey:uuid];
+        [NSURLProtocol setProperty:uuid forKey:key inRequest:request];
+    } else {
+        [NSURLProtocol setProperty:property forKey:key inRequest:request];
+    }
+}
+
++ (id)propertyForKey:(NSString *)key inRequest:(NSURLRequest *)request;
+{
+    id property = [NSURLProtocol propertyForKey:key inRequest:request];
+    return [storage objectForKey:property] ?: property;
+}
+
+@end

Sources/SBTUITestTunnelCommon/include/SBTRequestPropertyStorage.h

@@ -0,0 +1,22 @@
+//
+//  SBTRequestPropertyStorage.h
+//  Pods
+//
+//  Created by tomas on 20/02/24.
+//
+
+/// This class serves as a wrapper for NSURLProtocol to handle proxy property storage, addressing a limitation of 
+/// NSURLProtocol which restricts property size to 2^14 bytes. This class stores properties in a separate in memory
+/// storage assigning a unique uuid. This uuid is passed to the underlying NSURLProtocol and is used to retrieve the
+/// property from the storage when needed. Properties are stored in an NSDictionary with the uuid as the key and the
+/// property as the value. If this proves to be excessively optimistic memory wise an on disk storage can be implemented
+/// in the future.
+
+@import Foundation;
+
+@interface SBTRequestPropertyStorage : NSObject
+
++ (void)setProperty:(nonnull id)property forKey:(nonnull NSString *)key inRequest:(nonnull NSMutableURLRequest *)request;
++ (nullable id)propertyForKey:(nonnull NSString *)key inRequest:(nonnull NSURLRequest *)request;
+
+@end

Sources/SBTUITestTunnelServer/private/NSURLSession+HTTPBodyFix.m

@@ -26,7 +26,7 @@ - (NSURLSessionUploadTask *)swz_uploadTaskWithRequest:(NSURLRequest *)request fr
     NSURLRequest *requestWithoutBody = [request sbt_copyWithoutBody];
 
     if ([requestWithoutBody isKindOfClass:[NSMutableURLRequest class]] && bodyData) {
-        [NSURLProtocol setProperty:bodyData forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)requestWithoutBody];
+        [SBTRequestPropertyStorage setProperty:bodyData forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)requestWithoutBody];
 
         // mark this as an upload request so future code knows to find the body via NSURLProtocol instead
         [requestWithoutBody sbt_markUploadTaskRequest];
@@ -43,7 +43,7 @@ - (NSURLSessionUploadTask *)swz_uploadTaskWithRequest:(NSURLRequest *)request fr
     if ([requestWithoutBody isKindOfClass:[NSMutableURLRequest class]]) {
         NSData *bodyData = [NSData dataWithContentsOfURL:fileURL];
         if (bodyData) {
-            [NSURLProtocol setProperty:bodyData forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)requestWithoutBody];
+            [SBTRequestPropertyStorage setProperty:bodyData forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)requestWithoutBody];
 
             // mark this as an upload request so future code knows to find the body via NSURLProtocol instead
             [requestWithoutBody sbt_markUploadTaskRequest];
@@ -59,7 +59,7 @@ - (NSURLSessionUploadTask *)swz_uploadTaskWithRequest:(NSURLRequest *)request fr
     NSURLRequest *requestWithoutBody = [request sbt_copyWithoutBody];
 
     if ([requestWithoutBody isKindOfClass:[NSMutableURLRequest class]] && bodyData) {
-        [NSURLProtocol setProperty:bodyData forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)requestWithoutBody];
+        [SBTRequestPropertyStorage setProperty:bodyData forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)requestWithoutBody];
 
         // mark this as an upload request so future code knows to find the body via NSURLProtocol instead
         [requestWithoutBody sbt_markUploadTaskRequest];
@@ -76,7 +76,7 @@ - (NSURLSessionUploadTask *)swz_uploadTaskWithRequest:(NSURLRequest *)request fr
 - (NSURLSessionDataTask *)swz_dataTaskWithRequest:(NSURLRequest *)request
 {
     if ([request isKindOfClass:[NSMutableURLRequest class]] && request.HTTPBody) {
-        [NSURLProtocol setProperty:request.HTTPBody forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)request];
+        [SBTRequestPropertyStorage setProperty:request.HTTPBody forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)request];
     }
 
     return [self swz_dataTaskWithRequest:request];
@@ -85,7 +85,7 @@ - (NSURLSessionDataTask *)swz_dataTaskWithRequest:(NSURLRequest *)request
 - (NSURLSessionDataTask *)swz_dataTaskWithRequest:(NSURLRequest *)request completionHandler:(void (^)(NSData *data, NSURLResponse *response, NSError *error))completionHandler
 {
     if ([request isKindOfClass:[NSMutableURLRequest class]] && request.HTTPBody) {
-        [NSURLProtocol setProperty:request.HTTPBody forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)request];
+        [SBTRequestPropertyStorage setProperty:request.HTTPBody forKey:SBTUITunneledNSURLProtocolHTTPBodyKey inRequest:(NSMutableURLRequest *)request];
     }
 
     return [self swz_dataTaskWithRequest:request completionHandler:completionHandler];

Sources/SBTUITestTunnelServer/private/SBTProxyURLProtocol.m

@@ -362,7 +362,7 @@ + (BOOL)canInitWithRequest:(NSURLRequest *)request
     // values for the allHTTPHeaderFields property in one of these callse. For this reason
     // we postpone matching the request headers after startLoading is called.
 
-    if ([NSURLProtocol propertyForKey:SBTProxyURLProtocolHandledKey inRequest:request]) {
+    if ([SBTRequestPropertyStorage propertyForKey:SBTProxyURLProtocolHandledKey inRequest:request]) {
         return NO;
     }
 
@@ -449,7 +449,7 @@ - (void)startLoading
                     NSMutableURLRequest *redirectionRequest = [NSMutableURLRequest requestWithURL:redirectionUrl];
 
                     [NSURLProtocol removePropertyForKey:SBTProxyURLProtocolHandledKey inRequest:redirectionRequest];
-                    if (![NSURLProtocol propertyForKey:SBTProxyURLOriginalRequestKey inRequest:redirectionRequest]) {
+                    if (![SBTRequestPropertyStorage propertyForKey:SBTProxyURLOriginalRequestKey inRequest:redirectionRequest]) {
                         // don't handle double (or more) redirects
                         [[self class] associateOriginalRequest:request withRequest:redirectionRequest];
                     }
@@ -481,7 +481,7 @@ - (void)startLoading
         NSLog(@"[SBTUITestTunnel] Throttling/monitoring/chaning cookies/stubbing headers %@ request: %@\n\nMatching rule:\n%@", [self.request HTTPMethod], [self.request URL], requestMatch1 ?: requestMatch2 ?: requestMatch3 ?: requestMatch4 ?: requestMatch5);
 
         NSMutableURLRequest *newRequest = [self.request mutableCopy];
-        [NSURLProtocol setProperty:@YES forKey:SBTProxyURLProtocolHandledKey inRequest:newRequest];
+        [SBTRequestPropertyStorage setProperty:@YES forKey:SBTProxyURLProtocolHandledKey inRequest:newRequest];
 
         NSURLSession *session = [NSURLSession sessionWithConfiguration:[NSURLSessionConfiguration defaultSessionConfiguration] delegate:self delegateQueue:nil];
 
@@ -628,7 +628,7 @@ - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task willPer
     }
 
     [NSURLProtocol removePropertyForKey:SBTProxyURLProtocolHandledKey inRequest:mRequest];
-    if (![NSURLProtocol propertyForKey:SBTProxyURLOriginalRequestKey inRequest:mRequest]) {
+    if (![SBTRequestPropertyStorage propertyForKey:SBTProxyURLOriginalRequestKey inRequest:mRequest]) {
         // don't handle double (or more) redirects
         [[self class] associateOriginalRequest:self.request withRequest:mRequest];
     }
@@ -814,7 +814,7 @@ - (NSDictionary *)blockCookieRuleFromMatchingRules:(NSArray<NSDictionary *> *)ma
 
 /// Finds the original request in NSURLProtocol and deserializes it
 + (NSURLRequest *)originalRequestFor:(NSURLRequest*)request {
-    NSData *serializedOriginal = [NSURLProtocol propertyForKey:SBTProxyURLOriginalRequestKey inRequest:request];
+    NSData *serializedOriginal = [SBTRequestPropertyStorage propertyForKey:SBTProxyURLOriginalRequestKey inRequest:request];
     NSURLRequest *originalRequest = nil;
 
     if (serializedOriginal) {
@@ -837,7 +837,7 @@ + (void)associateOriginalRequest:(NSURLRequest *)original withRequest:(NSMutable
                                                                        error:&archiveError];
     NSAssert(archiveError == nil, @"Error archiving NSURLRequest for NSURLProtocol");
 
-    [NSURLProtocol setProperty:serializedOriginal forKey:SBTProxyURLOriginalRequestKey inRequest:request];
+    [SBTRequestPropertyStorage setProperty:serializedOriginal forKey:SBTProxyURLOriginalRequestKey inRequest:request];
 }
 
 @end