From 4bf4d2ba2b1ef00eb901ab121fa693a0e9dd5d3c Mon Sep 17 00:00:00 2001 From: Tomas Camin Date: Wed, 21 Feb 2024 09:25:51 +0100 Subject: [PATCH] Replace URLProtocol property setting --- .../SBTUITestTunnel_Tests/MonitorTests.swift | 36 ++++++++++++++++-- .../NSURLRequest+HTTPBodyFix.m | 17 +++++---- .../SBTRequestPropertyStorage.m | 38 +++++++++++++++++++ .../include/SBTRequestPropertyStorage.h | 22 +++++++++++ .../private/NSURLSession+HTTPBodyFix.m | 10 ++--- .../private/SBTProxyURLProtocol.m | 12 +++--- 6 files changed, 113 insertions(+), 22 deletions(-) create mode 100644 Sources/SBTUITestTunnelCommon/SBTRequestPropertyStorage.m create mode 100644 Sources/SBTUITestTunnelCommon/include/SBTRequestPropertyStorage.h diff --git a/Example/SBTUITestTunnel_Tests/MonitorTests.swift b/Example/SBTUITestTunnel_Tests/MonitorTests.swift index 64c93228..6a73894f 100644 --- a/Example/SBTUITestTunnel_Tests/MonitorTests.swift +++ b/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: "¶m5=val5¶m6=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), "¶m5=val5¶m6=val6") + XCTAssertEqual(String(data: httpBody, encoding: .utf8), largeBody) XCTAssert((request.responseString()!).contains("httpbin.org")) XCTAssert(request.timestamp > 0.0) diff --git a/Sources/SBTUITestTunnelCommon/NSURLRequest+HTTPBodyFix.m b/Sources/SBTUITestTunnelCommon/NSURLRequest+HTTPBodyFix.m index 1fd6da96..2c880ddb 100644 --- a/Sources/SBTUITestTunnelCommon/NSURLRequest+HTTPBodyFix.m +++ b/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,12 +38,12 @@ @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 @@ -50,7 +51,7 @@ - (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,7 +82,7 @@ - (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 @@ -89,9 +90,9 @@ - (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]; } } diff --git a/Sources/SBTUITestTunnelCommon/SBTRequestPropertyStorage.m b/Sources/SBTUITestTunnelCommon/SBTRequestPropertyStorage.m new file mode 100644 index 00000000..fb669253 --- /dev/null +++ b/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 diff --git a/Sources/SBTUITestTunnelCommon/include/SBTRequestPropertyStorage.h b/Sources/SBTUITestTunnelCommon/include/SBTRequestPropertyStorage.h new file mode 100644 index 00000000..e231f36b --- /dev/null +++ b/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 diff --git a/Sources/SBTUITestTunnelServer/private/NSURLSession+HTTPBodyFix.m b/Sources/SBTUITestTunnelServer/private/NSURLSession+HTTPBodyFix.m index 34cf0844..7b0d09ba 100644 --- a/Sources/SBTUITestTunnelServer/private/NSURLSession+HTTPBodyFix.m +++ b/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]; diff --git a/Sources/SBTUITestTunnelServer/private/SBTProxyURLProtocol.m b/Sources/SBTUITestTunnelServer/private/SBTProxyURLProtocol.m index 603a12c8..783ea2a2 100644 --- a/Sources/SBTUITestTunnelServer/private/SBTProxyURLProtocol.m +++ b/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 *)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