Attractor/App/Web Process Bundle Bridge/SBRProcessBundleBridge.m

//
//  SBRProcessBundleBridge.m
//  SBrowser
//
//  Created by James Magahern on 7/22/20.
//

#import "SBRProcessBundleBridge.h"

#import "SBRScriptPolicy.h"

#import <OSLog/OSLog.h>

#import <WebKit/_WKRemoteObjectInterface.h>
#import <WebKit/_WKRemoteObjectRegistry.h>
#import <WebKit/_WKProcessPoolConfiguration.h>
#import <WebKit/_WKUserStyleSheet.h>

#import <WebKit/WKProcessPoolPrivate.h>
#import <WebKit/WKWebViewPrivate.h>
#import <WebKit/WKWebViewConfigurationPrivate.h>
#import <WebKit/WKUserContentControllerPrivate.h>

#define LOG_DEBUG(format, ...) os_log_debug(_log, format, ##__VA_ARGS__)
#define LOG_ERROR(format, ...) os_log_error(_log, format, ##__VA_ARGS__)

@interface NSURLResponse (BridgeAdditions)
@property (nonatomic, readonly) BOOL isJavascriptResponse;
@end

@implementation NSURLResponse (BridgeAdditions)

- (BOOL)isJavascriptResponse
{
    NSString *extension = [[self URL] pathExtension];
    if ([[extension lowercaseString] isEqualToString:@"js"]) {
        return YES;
    }
    
    NSString *MIMEType = [self MIMEType];
    if ([[MIMEType lowercaseString] containsString:@"javascript"]) {
        return YES;
    }
    
    return NO;
}

@end

@interface SBRProcessBundleBridge () <WKURLSchemeHandler>

@end

@implementation SBRProcessBundleBridge {
    os_log_t                _log;
    
    WKWebView              *_webView;
    WKWebViewConfiguration *_webViewConfiguration;
    WKProcessPool          *_processPool;
    
    _WKUserStyleSheet     *_darkModeStyleSheet;
    WKUserScript          *_readabilityScript;
    
    NSArray<WKUserScript *> *_userScripts;
    
    dispatch_queue_t _dataTasksAccessQueue;
    NSMutableDictionary<NSURLRequest *, NSURLSessionDataTask *> *_dataTasks;
    
    // These come from settings. 
    _WKUserStyleSheet *_customizedUserStylesheet;
    WKUserScript      *_customizedUserScript;
}

- (void)tearDown
{
    // This was used to unregister the delegate with the web process.
}

- (instancetype)initWithWebViewConfiguration:(WKWebViewConfiguration *)webViewConfiguration
{
    self = [super init];
    if (self) {
        if (!webViewConfiguration) {
            _log = os_log_create("net.buzzert.attractor.webview", "bridge");
            
            webViewConfiguration = [[WKWebViewConfiguration alloc] init];

            // Set up process pool
            _processPool = [[WKProcessPool alloc] init];
            webViewConfiguration.processPool = _processPool;
            
            webViewConfiguration._waitsForPaintAfterViewDidMoveToWindow = NO;
            webViewConfiguration._applePayEnabled = YES;
            
            _dataTasks = [NSMutableDictionary dictionary];
            _dataTasksAccessQueue = dispatch_queue_create("net.buzzert.attractor.dataTasksAccess", DISPATCH_QUEUE_SERIAL);
            
            [webViewConfiguration setURLSchemeHandler:self forURLScheme:@"http"];
            [webViewConfiguration setURLSchemeHandler:self forURLScheme:@"https"];
        }
        
        _webViewConfiguration = webViewConfiguration;
        
        // User scripts
        WKUserContentController *userContentController = [_webViewConfiguration userContentController];
        for (WKUserScript *script in [self _userScripts]) {
            [userContentController addUserScript:script];
        }
        
        // Reload customized user scripts/stylesheets from settings
        [self reloadCustomizedUserScriptsAndStylesheets];
        
        // Instantiate web view
        WKWebView *webView = [[WKWebView alloc] initWithFrame:CGRectZero configuration:webViewConfiguration];
        if (@available(iOS 16.0, *)) {
            webView.findInteractionEnabled = YES;
        }
        
        _webView = webView;
    }
    
    return self;
}

- (WKUserScript *)_loadScriptForResource:(NSString *)resourceName withExtension:(NSString *)extension
{
    NSURL *url = [[NSBundle mainBundle] URLForResource:resourceName withExtension:extension];
    NSString *source = [NSString stringWithContentsOfURL:url encoding:NSUTF8StringEncoding error:nil];
    
    return [[WKUserScript alloc] initWithSource:source injectionTime:WKUserScriptInjectionTimeAtDocumentEnd forMainFrameOnly:YES];
}

- (NSArray<WKUserScript *> *)_userScripts
{
    if (!_userScripts) {
        _userScripts = @[
            [self _loadScriptForResource:@"Tagger" withExtension:@"js"],
        ];
    }
    
    return _userScripts;
}

- (void)reloadCustomizedUserScriptsAndStylesheets
{
    WKUserContentController *userContentController = [_webViewConfiguration userContentController];
    if (_customizedUserScript) {
        [userContentController _removeUserScript:_customizedUserScript];
    }
    
    NSString *scriptSource = [[NSUserDefaults standardUserDefaults] stringForKey:@"userScript"];
    if ([scriptSource length]) {
        _customizedUserScript = [[WKUserScript alloc] initWithSource:scriptSource injectionTime:WKUserScriptInjectionTimeAtDocumentEnd forMainFrameOnly:YES];
        [userContentController addUserScript:_customizedUserScript];
    }
    
    if (_customizedUserStylesheet) {
        [userContentController _removeUserStyleSheet:_customizedUserStylesheet];
    }
    
    NSString *stylesheetSource = [[NSUserDefaults standardUserDefaults] stringForKey:@"userStylesheet"];
    if ([stylesheetSource length]) {
        _customizedUserStylesheet = [[_WKUserStyleSheet alloc] initWithSource:stylesheetSource forMainFrameOnly:YES];
        [userContentController _addUserStyleSheet:_customizedUserStylesheet];
    }
}

#pragma mark Former <SBRWebProcessDelegate> methods

- (void)webProcessDidAllowScriptWithOrigin:(NSString *)origin
{
    dispatch_async(dispatch_get_main_queue(), ^{
        [[self delegate] webProcess:self didAllowScriptResourceFromOrigin:origin];
    });
}

- (void)webProcessDidBlockScriptWithOrigin:(NSString *)origin
{
    dispatch_async(dispatch_get_main_queue(), ^{
        [[self delegate] webProcess:self didBlockScriptResourceFromOrigin:origin];
    });
}

#pragma mark <WKURLSchemeHandler>

- (void)webView:(WKWebView *)webView startURLSchemeTask:(id<WKURLSchemeTask>)urlSchemeTask
{
    NSString *hostOrigin = [[_webView URL] host];
    NSURLRequest *request = [urlSchemeTask request];
    
    LOG_DEBUG("Start URL scheme task: request: %@", request);
    
    __weak __auto_type welf = self;
    NSURLSessionDataTask *dataTask = [[NSURLSession sharedSession] dataTaskWithRequest:request completionHandler:^(NSData * _Nullable data, NSURLResponse * _Nullable response, NSError * _Nullable error)
    {
        if (!welf) return;
        __strong __auto_type sself = welf;
        
        if (error != nil) {
            [urlSchemeTask didFailWithError:error];
        } else if ([response isKindOfClass:[NSHTTPURLResponse class]]) {
            NSURL *requestURL = [request URL];
            NSString *resourceOrigin = [requestURL host];
            const __auto_type allowResource = ^{
                os_log_debug(sself->_log, "Allowing resource: %@", requestURL.lastPathComponent);
                [urlSchemeTask didReceiveResponse:response];
                [urlSchemeTask didReceiveData:data];
                [urlSchemeTask didFinish];
                
                [self webProcessDidAllowScriptWithOrigin:resourceOrigin];
            };
            
            const __auto_type denyResource = ^{
                os_log_debug(sself->_log, "Blocking resource: %@", requestURL.lastPathComponent);
                NSHTTPURLResponse *altResponse = [[NSHTTPURLResponse alloc] initWithURL:requestURL
                                                                               MIMEType:@"application/javascript"
                                                                  expectedContentLength:0 textEncodingName:@"utf8"];
                [urlSchemeTask didReceiveResponse:altResponse];
                [urlSchemeTask didReceiveData:[NSData data]];
                [urlSchemeTask didFinish];
                
                [self webProcessDidBlockScriptWithOrigin:resourceOrigin];
            };
            
            // Check MIME type for JavaScript responses.
            if ([response isJavascriptResponse] && ![sself allowAllScripts]) {
                dispatch_async(sself->_dataTasksAccessQueue, ^{
                    NSDictionary<NSString *, NSNumber *> *policyTypes = [sself->_policyDataSource scriptPolicyTypeByOrigin];
                    NSNumber *policyType = [policyTypes objectForKey:hostOrigin];
                    
                    SBRScriptPolicy *policy = [[SBRScriptPolicy alloc] initWithSecurityOrigin:hostOrigin policyType:[policyType integerValue]];
                    if ([policy allowsExternalJavaScriptResourceOrigin:resourceOrigin]) {
                        allowResource();
                    } else {
                        denyResource();
                    }
                });
            } else {
                allowResource();
            }
        } else {
            [urlSchemeTask didFailWithError:[NSError errorWithDomain:NSURLErrorDomain code:0 userInfo:nil]];
        }
        
        [sself->_dataTasks removeObjectForKey:request];
    }];
    
    [_dataTasks setObject:dataTask forKey:request];
    [dataTask resume];
}

- (void)webView:(WKWebView *)webView stopURLSchemeTask:(id<WKURLSchemeTask>)urlSchemeTask
{
    NSURLRequest *request = [urlSchemeTask request];
    NSURLSessionDataTask *dataTask = [_dataTasks objectForKey:request];
    if (dataTask) {
        if ([dataTask state] != NSURLSessionTaskStateCanceling) {
            [dataTask cancel];
        }
        
        [_dataTasks removeObjectForKey:request];
    }
}

#pragma mark Actions

- (void)policyDataSourceDidChange
{
    // This was used when we had to signal the process bundle.
}

- (void)setAllowAllScripts:(BOOL)allowAllScripts
{
    _allowAllScripts = allowAllScripts;
}

- (void)setDarkModeEnabled:(BOOL)darkModeEnabled
{
    _darkModeEnabled = darkModeEnabled;
    
    WKUserContentController *userContentController = [_webViewConfiguration userContentController];
    
    if (darkModeEnabled) {
        if (!_darkModeStyleSheet) {
            NSURL *styleSheetURL = [[NSBundle mainBundle] URLForResource:@"darkmode" withExtension:@"css"];
            NSString *styleSheetSource = [NSString stringWithContentsOfURL:styleSheetURL encoding:NSUTF8StringEncoding error:nil];
            _darkModeStyleSheet = [[_WKUserStyleSheet alloc] initWithSource:styleSheetSource forMainFrameOnly:NO];
        }
        
        [userContentController _addUserStyleSheet:_darkModeStyleSheet];
    } else if (_darkModeStyleSheet) {
        [userContentController _removeUserStyleSheet:_darkModeStyleSheet];
    }
}

- (void)parseDocumentForReaderMode:(void (^)(NSString * _Nonnull))completionBlock
{
    WKUserContentController *userContentController = [_webViewConfiguration userContentController];
    
    if (!_readabilityScript) {
        _readabilityScript = [self _loadScriptForResource:@"Readability" withExtension:@"js"];
    }
    
    [userContentController _addUserScriptImmediately:_readabilityScript];
    
    NSString *script = @""
    "var documentClone = document.cloneNode(true);"
    "var article = new Readability(documentClone).parse();"
    "article.content";
    
    os_log_t log = _log;
    [_webView evaluateJavaScript:script completionHandler:^(NSString *result, NSError * _Nullable error) {
        if (error != nil) {
            os_log_error(log, "Bridge: Readability error: %@", error.localizedDescription);
        } else {
            completionBlock(result);
        }
    }];
}

@end