Attractor/SBrowserProcessBundle/SBRProcessPlugin.m

//
//  SBRProcessPlugin.m
//  SBrowserProcessBundle
//
//  Created by James Magahern on 7/22/20.
//

#import "SBRProcessPlugin.h"

#import "SBRWebProcessDelegate.h"
#import "SBRWebProcessProxy.h"
#import "SBRScriptPolicy.h"

#import <WebKit/_WKRemoteObjectInterface.h>
#import <WebKit/_WKRemoteObjectRegistry.h>
#import <WebKit/WKWebProcessPlugInBrowserContextController.h>
#import <WebKit/WKWebProcessPlugInBrowserContextControllerPrivate.h>
#import <WebKit/WKWebProcessPlugInFrame.h>
#import <WebKit/WKWebProcessPlugInLoadDelegate.h>

@interface SBRProcessPlugin () <WKWebProcessPlugInLoadDelegate, SBRWebProcessProxy>

@property (nonatomic, strong) id<SBRWebProcessDelegate> processDelegate;

@property (nonatomic, assign) BOOL allScriptsAllowed;
@property (nonatomic, strong) NSMutableSet<NSString *> *allowedResourceOrigins;
@property (nonatomic, strong) NSDictionary<NSString *, NSNumber *> *policyTypeByOrigin;
@end

@implementation SBRProcessPlugin

+ (void)initialize
{
    [super initialize];
    NSLog(@"SBRProcessPlugin Initialized");
}

#pragma mark <SBRWebProcessProxy>

- (void)hello
{
    NSLog(@"SBRProcessPlugin: Helloooooo");
}

- (void)syncAllowedResourceOrigins:(NSArray<NSString *> *)allowedOrigins
{
    _allowedResourceOrigins = [NSMutableSet setWithArray:allowedOrigins];
}

- (void)syncPolicyTypes:(NSDictionary<NSString *,NSNumber *> *)policyTypes
{
    _policyTypeByOrigin = policyTypes;
}

- (void)setAllScriptsAllowed:(BOOL)allScriptsAllowed
{
    _allScriptsAllowed = allScriptsAllowed;
}

#pragma mark <WKWebProcessPlugIn>

- (void)webProcessPlugIn:(WKWebProcessPlugInController *)plugInController initializeWithObject:(id)initializationObject
{
    _allowedResourceOrigins = [[plugInController parameters] valueForKey:SBRGetAllowedOriginsKey()];
    _allScriptsAllowed = [[[plugInController parameters] valueForKey:SBRGetAllScriptsAllowedKey()] boolValue];
    _policyTypeByOrigin = [[plugInController parameters] valueForKey:SBRGetPolicyTypeByOriginKey()];
    NSLog(@"SBRProcessPlugin: %lu origins allowed, all scripts allowed: %@ ", (unsigned long)_allowedResourceOrigins.count, _allScriptsAllowed ? @"YES" : @"NO");
}

- (void)webProcessPlugIn:(WKWebProcessPlugInController *)plugInController didCreateBrowserContextController:(WKWebProcessPlugInBrowserContextController *)browserContextController
{
    _WKRemoteObjectInterface *proxyInterface = [_WKRemoteObjectInterface remoteObjectInterfaceWithProtocol:@protocol(SBRWebProcessProxy)];
    [browserContextController._remoteObjectRegistry registerExportedObject:self interface:proxyInterface];

    _WKRemoteObjectInterface *eventListenerInterface = [_WKRemoteObjectInterface remoteObjectInterfaceWithProtocol:@protocol(SBRWebProcessDelegate)];
    self.processDelegate = [browserContextController._remoteObjectRegistry remoteObjectProxyWithInterface:eventListenerInterface];

    browserContextController.loadDelegate = self;
    
    [[self processDelegate] webProcessDidConnect];
}

#pragma mark <WKWebProcessPlugInLoadDelegate>

- (NSURLRequest *)webProcessPlugInBrowserContextController:(WKWebProcessPlugInBrowserContextController *)controller frame:(WKWebProcessPlugInFrame *)frame willSendRequest:(NSURLRequest *)request redirectResponse:(NSURLResponse *)redirectResponse
{
    NSLog(@"SBRProcessPlugin: Sending request: %@", request);
    if (_allowedResourceOrigins == nil) {
        NSLog(@"Allowed resource origins should not be nil!!!!");
    }
    
    NSURL *requestURL = [request URL];
    NSString *resourceOrigin = [requestURL host];
    NSString *requestExtension = [requestURL pathExtension];
    NSString *hostOrigin = [[[controller mainFrame] URL] host];
    if (requestExtension.length > 0 && [requestExtension isEqualToString:@"js"]) {
        NSNumber *policyType = [_policyTypeByOrigin objectForKey:hostOrigin];
        NSLog(@"SBRProcessPlugin: Policy type for %@: %@", hostOrigin, policyType);
        
        SBRScriptPolicy *policy = [[SBRScriptPolicy alloc] initWithSecurityOrigin:hostOrigin policyType:[policyType integerValue]];
        if ([self allScriptsAllowed] || [policy allowsExternalJavaScriptResourceOrigin:resourceOrigin]) {
            NSLog(@"SBRProcessPlugin: Policy allows script requestURL: %@", requestURL);
            [[self processDelegate] webProcessDidAllowScriptWithOrigin:resourceOrigin];
        } else {
            NSLog(@"SBRProcessPlugin: Policy disallows script requestURL: %@", requestURL);
            [[self processDelegate] webProcessDidBlockScriptWithOrigin:resourceOrigin];
            
            request = nil;
        }
    }
    
    return request;
}

@end