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

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

#import "SBRProcessBundleBridge.h"

#import "SBRWebProcessDelegate.h"
#import "SBRWebProcessProxy.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>

@interface SBRProcessBundleBridge () <SBRWebProcessDelegate>

@end

@implementation SBRProcessBundleBridge {
    WKWebView              *_webView;
    WKWebViewConfiguration *_webViewConfiguration;
    WKProcessPool          *_processPool;
    id<SBRWebProcessProxy> _webProcessProxy;
    
    _WKUserStyleSheet     *_darkModeStyleSheet;
    WKUserScript          *_readabilityScript;
    
    NSArray<WKUserScript *> *_userScripts;
    
    // These come from settings. 
    _WKUserStyleSheet *_customizedUserStylesheet;
    WKUserScript      *_customizedUserScript;
}

- (void)tearDown
{
    [[_webView _remoteObjectRegistry] unregisterExportedObject:self interface:[self _webProcessDelegateInterface]];
}

- (_WKRemoteObjectInterface *)_webProcessDelegateInterface
{
    static dispatch_once_t onceToken;
    static _WKRemoteObjectInterface *interface = nil;
    dispatch_once(&onceToken, ^{
        interface = [_WKRemoteObjectInterface remoteObjectInterfaceWithProtocol:@protocol(SBRWebProcessDelegate)];
    });
    
    return interface;
}

- (_WKRemoteObjectInterface *)_webProcessProxyInterface
{
    static dispatch_once_t onceToken;
    static _WKRemoteObjectInterface *interface = nil;
    dispatch_once(&onceToken, ^{
        interface = [_WKRemoteObjectInterface remoteObjectInterfaceWithProtocol:@protocol(SBRWebProcessProxy)];
    });
    
    return interface;
}

- (instancetype)initWithWebViewConfiguration:(WKWebViewConfiguration *)webViewConfiguration
{
    self = [super init];
    if (self) {
        if (!webViewConfiguration) {
            webViewConfiguration = [[WKWebViewConfiguration alloc] init];
            
            // Inject bundle
            _WKProcessPoolConfiguration *poolConfiguration = [[_WKProcessPoolConfiguration alloc] init];
            NSURL *bundleURL = [[[NSBundle mainBundle] builtInPlugInsURL] URLByAppendingPathComponent:@"SBrowserProcessBundle.bundle"];
            
            // Make sure it exists. Bail if otherwise.
            NSBundle *pluginBundle = [NSBundle bundleWithURL:bundleURL];
            NSAssert(pluginBundle != nil, @"Attix process bundle not found at %@", bundleURL.path);
            
            [poolConfiguration setInjectedBundleURL:bundleURL];

            // Set up process pool
            _processPool = [[WKProcessPool alloc] _initWithConfiguration:poolConfiguration];
            
            // Initialize allowed origins now
            NSDictionary<NSString *, NSNumber *> *policies = [_policyDataSource scriptPolicyTypeByOrigin];
            NSArray<NSString *> *allowedOrigins = [[_policyDataSource allowedOriginsForScriptResources] allObjects];
            [_processPool _setObject:allowedOrigins forBundleParameter:SBRGetAllowedOriginsKey()];
            [_processPool _setObject:@(_allowAllScripts) forBundleParameter:SBRGetAllScriptsAllowedKey()];
            [_processPool _setObject:policies forBundleParameter:SBRGetPolicyTypeByOriginKey()];
            
            webViewConfiguration.processPool = _processPool;
            
            webViewConfiguration._waitsForPaintAfterViewDidMoveToWindow = NO;
            webViewConfiguration._applePayEnabled = YES;
        }
        
        _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;
        }
        
        // Configure proxy interface (interface to remote web process)
        _webProcessProxy = [[webView _remoteObjectRegistry] remoteObjectProxyWithInterface:[self _webProcessProxyInterface]];
        
        // Configure delegate interface (registering us as the web process delegate for the remote process)
        [[webView _remoteObjectRegistry] registerExportedObject:self interface:[self _webProcessDelegateInterface]];
        
        _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 <SBRWebProcessDelegate>

- (void)webProcessDidConnect
{
    NSLog(@"SBRProcessBundleBridge: did connect. Saying hello, syncing allowlist");
    _webContentProcessConnected = YES;
    
    [_webProcessProxy hello];
    [self policyDataSourceDidChange];
}

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

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

#pragma mark Actions

- (void)policyDataSourceDidChange
{
    NSArray<NSString *> *allowedOrigins = [[_policyDataSource allowedOriginsForScriptResources] allObjects];
    [_processPool _setObject:allowedOrigins forBundleParameter:SBRGetAllowedOriginsKey()];
    [_webProcessProxy syncAllowedResourceOrigins:allowedOrigins];
    
    NSDictionary<NSString *, NSNumber *> *policyTypes = [_policyDataSource scriptPolicyTypeByOrigin];
    [_processPool _setObject:policyTypes forBundleParameter:SBRGetPolicyTypeByOriginKey()];
    [_webProcessProxy syncPolicyTypes:policyTypes];
}

- (void)setAllowAllScripts:(BOOL)allowAllScripts
{
    _allowAllScripts = allowAllScripts;
    [_processPool _setObject:@(_allowAllScripts) forBundleParameter:SBRGetAllScriptsAllowedKey()];
    [_webProcessProxy setAllScriptsAllowed: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";
    
    [_webView evaluateJavaScript:script completionHandler:^(NSString *result, NSError * _Nullable error) {
        if (error != nil) {
            NSLog(@"Bridge: Readability error: %@", error.localizedDescription);
        } else {
            completionBlock(result);
        }
    }];
}

@end