Skip to content

Commit

Permalink
Hook and record websocket methods
Browse files Browse the repository at this point in the history
  • Loading branch information
@NSExceptional
NSExceptional committed Oct 6, 2021
1 parent ad1f1f5 commit 3446eff
Show file tree
Hide file tree
Showing 3 changed files with 220 additions and 58 deletions.
22 changes: 16 additions & 6 deletions Classes/Network/FLEXNetworkRecorder.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ extern NSString *const kFLEXNetworkRecorderTransactionUpdatedNotification;
extern NSString *const kFLEXNetworkRecorderUserInfoTransactionKey;
extern NSString *const kFLEXNetworkRecorderTransactionsClearedNotification;

@class FLEXNetworkTransaction;
@class FLEXNetworkTransaction, FLEXHTTPTransaction, FLEXWebsocketTransaction;

@interface FLEXNetworkRecorder : NSObject

Expand All @@ -37,19 +37,21 @@ extern NSString *const kFLEXNetworkRecorderTransactionsClearedNotification;
- (void)synchronizeDenylist;


// Accessing recorded network activity
#pragma mark Accessing recorded network activity

/// Array of FLEXNetworkTransaction objects ordered by start time with the newest first.
- (NSArray<FLEXNetworkTransaction *> *)networkTransactions;
/// Array of FLEXHTTPTransaction objects ordered by start time with the newest first.
@property (nonatomic, readonly) NSArray<FLEXHTTPTransaction *> *HTTPTransactions;
/// Array of FLEXWebsocketTransaction objects ordered by start time with the newest first.
@property (nonatomic, readonly) NSArray<FLEXWebsocketTransaction *> *websocketTransactions API_AVAILABLE(ios(13.0));

/// The full response data IFF it hasn't been purged due to memory pressure.
- (NSData *)cachedResponseBodyForTransaction:(FLEXNetworkTransaction *)transaction;
- (NSData *)cachedResponseBodyForTransaction:(FLEXHTTPTransaction *)transaction;

/// Dumps all network transactions and cached response bodies.
- (void)clearRecordedActivity;


// Recording network activity
#pragma mark Recording network activity

/// Call when app is about to send HTTP request.
- (void)recordRequestWillBeSentWithRequestID:(NSString *)requestID
Expand All @@ -72,4 +74,12 @@ extern NSString *const kFLEXNetworkRecorderTransactionsClearedNotification;
/// This string can be set to anything useful about the API used to make the request.
- (void)recordMechanism:(NSString *)mechanism forRequestID:(NSString *)requestID;

- (void)recordWebsocketMessageSend:(NSURLSessionWebSocketMessage *)message
task:(NSURLSessionWebSocketTask *)task API_AVAILABLE(ios(13.0));
- (void)recordWebsocketMessageSendCompletion:(NSURLSessionWebSocketMessage *)message
error:(NSError *)error API_AVAILABLE(ios(13.0));

- (void)recordWebsocketMessageReceived:(NSURLSessionWebSocketMessage *)message
task:(NSURLSessionWebSocketTask *)task API_AVAILABLE(ios(13.0));

@end
132 changes: 83 additions & 49 deletions Classes/Network/FLEXNetworkRecorder.m
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
#import "FLEXUtility.h"
#import "FLEXResources.h"
#import "NSUserDefaults+FLEX.h"
#import "OSCache.h"

NSString *const kFLEXNetworkRecorderNewTransactionNotification = @"kFLEXNetworkRecorderNewTransactionNotification";
NSString *const kFLEXNetworkRecorderTransactionUpdatedNotification = @"kFLEXNetworkRecorderTransactionUpdatedNotification";
Expand All @@ -22,9 +23,10 @@

@interface FLEXNetworkRecorder ()

@property (nonatomic) NSCache *responseCache;
@property (nonatomic) NSMutableArray<FLEXNetworkTransaction *> *orderedTransactions;
@property (nonatomic) NSMutableDictionary<NSString *, FLEXNetworkTransaction *> *requestIDsToTransactions;
@property (nonatomic) OSCache *restCache;
@property (nonatomic) NSMutableArray<FLEXHTTPTransaction *> *orderedHTTPTransactions;
@property (nonatomic) NSMutableArray<FLEXWebsocketTransaction *> *orderedWSTransactions;
@property (nonatomic) NSMutableDictionary<NSString *, FLEXHTTPTransaction *> *requestIDsToHTTPTransactions;
@property (nonatomic) dispatch_queue_t queue;

@end
Expand All @@ -34,17 +36,18 @@ @implementation FLEXNetworkRecorder
- (instancetype)init {
self = [super init];
if (self) {
self.responseCache = [NSCache new];
self.restCache = [OSCache new];
NSUInteger responseCacheLimit = [[NSUserDefaults.standardUserDefaults
objectForKey:kFLEXNetworkRecorderResponseCacheLimitDefaultsKey] unsignedIntegerValue
];

// Default to 25 MB max. The cache will purge earlier if there is memory pressure.
self.responseCache.totalCostLimit = responseCacheLimit ?: 25 * 1024 * 1024;
[self.responseCache setTotalCostLimit:responseCacheLimit];
self.restCache.totalCostLimit = responseCacheLimit ?: 25 * 1024 * 1024;
[self.restCache setTotalCostLimit:responseCacheLimit];

self.orderedTransactions = [NSMutableArray new];
self.requestIDsToTransactions = [NSMutableDictionary new];
self.orderedWSTransactions = [NSMutableArray new];
self.orderedHTTPTransactions = [NSMutableArray new];
self.requestIDsToHTTPTransactions = [NSMutableDictionary new];
self.hostDenylist = NSUserDefaults.standardUserDefaults.flex_networkHostDenylist.mutableCopy;

// Serial queue used because we use mutable objects that are not thread safe
Expand All @@ -67,43 +70,43 @@ + (instancetype)defaultRecorder {
#pragma mark - Public Data Access

- (NSUInteger)responseCacheByteLimit {
return self.responseCache.totalCostLimit;
return self.restCache.totalCostLimit;
}

- (void)setResponseCacheByteLimit:(NSUInteger)responseCacheByteLimit {
self.responseCache.totalCostLimit = responseCacheByteLimit;
self.restCache.totalCostLimit = responseCacheByteLimit;
[NSUserDefaults.standardUserDefaults
setObject:@(responseCacheByteLimit)
forKey:kFLEXNetworkRecorderResponseCacheLimitDefaultsKey
];
}

- (NSArray<FLEXNetworkTransaction *> *)networkTransactions {
__block NSArray<FLEXNetworkTransaction *> *transactions = nil;
dispatch_sync(self.queue, ^{
transactions = self.orderedTransactions.copy;
});
return transactions;
- (NSArray<FLEXHTTPTransaction *> *)HTTPTransactions {
return self.orderedHTTPTransactions.copy;
}

- (NSArray<FLEXWebsocketTransaction *> *)websocketTransactions {
return self.orderedWSTransactions.copy;
}

- (NSData *)cachedResponseBodyForTransaction:(FLEXNetworkTransaction *)transaction {
return [self.responseCache objectForKey:transaction.requestID];
- (NSData *)cachedResponseBodyForTransaction:(FLEXHTTPTransaction *)transaction {
return [self.restCache objectForKey:transaction.requestID];
}

- (void)clearRecordedActivity {
dispatch_async(self.queue, ^{
[self.responseCache removeAllObjects];
[self.orderedTransactions removeAllObjects];
[self.requestIDsToTransactions removeAllObjects];
[self.restCache removeAllObjects];
[self.orderedHTTPTransactions removeAllObjects];
[self.requestIDsToHTTPTransactions removeAllObjects];

[self notify:kFLEXNetworkRecorderTransactionsClearedNotification transaction:nil];
});
}

- (void)clearExcludedTransactions {
dispatch_sync(self.queue, ^{
self.orderedTransactions = ({
[self.orderedTransactions flex_filtered:^BOOL(FLEXNetworkTransaction *ta, NSUInteger idx) {
self.orderedHTTPTransactions = ({
[self.orderedHTTPTransactions flex_filtered:^BOOL(FLEXHTTPTransaction *ta, NSUInteger idx) {
NSString *host = ta.request.URL.host;
for (NSString *excluded in self.hostDenylist) {
if ([host hasSuffix:excluded]) {
Expand Down Expand Up @@ -132,22 +135,17 @@ - (void)recordRequestWillBeSentWithRequestID:(NSString *)requestID
}
}

// Before async block to stay accurate
NSDate *startDate = [NSDate date];
FLEXHTTPTransaction *transaction = [FLEXHTTPTransaction request:request identifier:requestID];

// Before async block to keep times accurate
if (redirectResponse) {
[self recordResponseReceivedWithRequestID:requestID response:redirectResponse];
[self recordLoadingFinishedWithRequestID:requestID responseBody:nil];
}

dispatch_async(self.queue, ^{
FLEXNetworkTransaction *transaction = [FLEXNetworkTransaction new];
transaction.requestID = requestID;
transaction.request = request;
transaction.startTime = startDate;

[self.orderedTransactions insertObject:transaction atIndex:0];
[self.requestIDsToTransactions setObject:transaction forKey:requestID];
[self.orderedHTTPTransactions insertObject:transaction atIndex:0];
[self.requestIDsToHTTPTransactions setObject:transaction forKey:requestID];
transaction.transactionState = FLEXNetworkTransactionStateAwaitingResponse;

[self postNewTransactionNotificationWithTransaction:transaction];
Expand All @@ -159,7 +157,7 @@ - (void)recordResponseReceivedWithRequestID:(NSString *)requestID response:(NSUR
NSDate *responseDate = [NSDate date];

dispatch_async(self.queue, ^{
FLEXNetworkTransaction *transaction = self.requestIDsToTransactions[requestID];
FLEXHTTPTransaction *transaction = self.requestIDsToHTTPTransactions[requestID];
if (!transaction) {
return;
}
Expand All @@ -174,7 +172,7 @@ - (void)recordResponseReceivedWithRequestID:(NSString *)requestID response:(NSUR

- (void)recordDataReceivedWithRequestID:(NSString *)requestID dataLength:(int64_t)dataLength {
dispatch_async(self.queue, ^{
FLEXNetworkTransaction *transaction = self.requestIDsToTransactions[requestID];
FLEXHTTPTransaction *transaction = self.requestIDsToHTTPTransactions[requestID];
if (!transaction) {
return;
}
Expand All @@ -188,7 +186,7 @@ - (void)recordLoadingFinishedWithRequestID:(NSString *)requestID responseBody:(N
NSDate *finishedDate = [NSDate date];

dispatch_async(self.queue, ^{
FLEXNetworkTransaction *transaction = self.requestIDsToTransactions[requestID];
FLEXHTTPTransaction *transaction = self.requestIDsToHTTPTransactions[requestID];
if (!transaction) {
return;
}
Expand All @@ -205,40 +203,40 @@ - (void)recordLoadingFinishedWithRequestID:(NSString *)requestID responseBody:(N
}

if (shouldCache) {
[self.responseCache setObject:responseBody forKey:requestID cost:responseBody.length];
[self.restCache setObject:responseBody forKey:requestID cost:responseBody.length];
}

NSString *mimeType = transaction.response.MIMEType;
if ([mimeType hasPrefix:@"image/"] && responseBody.length > 0) {
// Thumbnail image previews on a separate background queue
dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
NSInteger maxPixelDimension = UIScreen.mainScreen.scale * 32.0;
transaction.responseThumbnail = [FLEXUtility
transaction.thumbnail = [FLEXUtility
thumbnailedImageWithMaxPixelDimension:maxPixelDimension
fromImageData:responseBody
];
[self postUpdateNotificationForTransaction:transaction];
});
} else if ([mimeType isEqual:@"application/json"]) {
transaction.responseThumbnail = FLEXResources.jsonIcon;
transaction.thumbnail = FLEXResources.jsonIcon;
} else if ([mimeType isEqual:@"text/plain"]){
transaction.responseThumbnail = FLEXResources.textPlainIcon;
transaction.thumbnail = FLEXResources.textPlainIcon;
} else if ([mimeType isEqual:@"text/html"]) {
transaction.responseThumbnail = FLEXResources.htmlIcon;
transaction.thumbnail = FLEXResources.htmlIcon;
} else if ([mimeType isEqual:@"application/x-plist"]) {
transaction.responseThumbnail = FLEXResources.plistIcon;
transaction.thumbnail = FLEXResources.plistIcon;
} else if ([mimeType isEqual:@"application/octet-stream"] || [mimeType isEqual:@"application/binary"]) {
transaction.responseThumbnail = FLEXResources.binaryIcon;
transaction.thumbnail = FLEXResources.binaryIcon;
} else if ([mimeType containsString:@"javascript"]) {
transaction.responseThumbnail = FLEXResources.jsIcon;
transaction.thumbnail = FLEXResources.jsIcon;
} else if ([mimeType containsString:@"xml"]) {
transaction.responseThumbnail = FLEXResources.xmlIcon;
transaction.thumbnail = FLEXResources.xmlIcon;
} else if ([mimeType hasPrefix:@"audio"]) {
transaction.responseThumbnail = FLEXResources.audioIcon;
transaction.thumbnail = FLEXResources.audioIcon;
} else if ([mimeType hasPrefix:@"video"]) {
transaction.responseThumbnail = FLEXResources.videoIcon;
transaction.thumbnail = FLEXResources.videoIcon;
} else if ([mimeType hasPrefix:@"text"]) {
transaction.responseThumbnail = FLEXResources.textIcon;
transaction.thumbnail = FLEXResources.textIcon;
}

[self postUpdateNotificationForTransaction:transaction];
Expand All @@ -247,7 +245,7 @@ - (void)recordLoadingFinishedWithRequestID:(NSString *)requestID responseBody:(N

- (void)recordLoadingFailedWithRequestID:(NSString *)requestID error:(NSError *)error {
dispatch_async(self.queue, ^{
FLEXNetworkTransaction *transaction = self.requestIDsToTransactions[requestID];
FLEXHTTPTransaction *transaction = self.requestIDsToHTTPTransactions[requestID];
if (!transaction) {
return;
}
Expand All @@ -262,7 +260,7 @@ - (void)recordLoadingFailedWithRequestID:(NSString *)requestID error:(NSError *)

- (void)recordMechanism:(NSString *)mechanism forRequestID:(NSString *)requestID {
dispatch_async(self.queue, ^{
FLEXNetworkTransaction *transaction = self.requestIDsToTransactions[requestID];
FLEXHTTPTransaction *transaction = self.requestIDsToHTTPTransactions[requestID];
if (!transaction) {
return;
}
Expand All @@ -272,6 +270,42 @@ - (void)recordMechanism:(NSString *)mechanism forRequestID:(NSString *)requestID
});
}

#pragma mark - Websocket Events

- (void)recordWebsocketMessageSend:(NSURLSessionWebSocketMessage *)message task:(NSURLSessionWebSocketTask *)task {
dispatch_async(self.queue, ^{
FLEXWebsocketTransaction *send = [FLEXWebsocketTransaction
withMessage:message task:task direction:FLEXWebsocketOutgoing
];

[self.orderedWSTransactions addObject:send];
[self postNewTransactionNotificationWithTransaction:send];
});
}

- (void)recordWebsocketMessageSendCompletion:(NSURLSessionWebSocketMessage *)message error:(NSError *)error {
dispatch_async(self.queue, ^{
FLEXWebsocketTransaction *send = [self.orderedWSTransactions flex_firstWhere:^BOOL(FLEXWebsocketTransaction *t) {
return t.message == message;
}];
send.error = error;
send.transactionState = error ? FLEXNetworkTransactionStateFailed : FLEXNetworkTransactionStateFinished;

[self postUpdateNotificationForTransaction:send];
});
}

- (void)recordWebsocketMessageReceived:(NSURLSessionWebSocketMessage *)message task:(NSURLSessionWebSocketTask *)task {
dispatch_async(self.queue, ^{
FLEXWebsocketTransaction *receive = [FLEXWebsocketTransaction
withMessage:message task:task direction:FLEXWebsocketIncoming
];

[self.orderedWSTransactions addObject:receive];
[self postNewTransactionNotificationWithTransaction:receive];
});
}

#pragma mark Notification Posting

- (void)postNewTransactionNotificationWithTransaction:(FLEXNetworkTransaction *)transaction {
Expand Down
Loading

0 comments on commit 3446eff

Please sign in to comment.