inject-scratch-link.js

class ScratchLink {

    constructor(url) {
        this.url = url;
        this._open();
    }

    _open() {
        this.socketId = ScratchLink.socketId;
        ScratchLink.sockets.set(ScratchLink.socketId, this);
        ScratchLink.socketId++;

        this._postMessage({
            method: 'open',
            socketId: this.socketId,
            url: this.url
        });

        setTimeout(() => {
            this.onopen();
        }, 100);
    }

    close() {
        this._postMessage({
            method: 'close',
            socketId: this.socketId
        })

        this.onclose();

        ScratchLink.sockets.delete(this.socketId);
    }

    send(message) {
        const data = JSON.parse(message);

        this._postMessage({
            method: 'send',
            socketId: this.socketId,
            jsonrpc: data
        });
    }

    _postMessage(message) {
        webkit.messageHandlers.callbackHandler.postMessage(JSON.stringify(message));
    }

    handleMessage(message) {
        this.onmessage({
            data: JSON.stringify(message)
        });
    }
}

ScratchLink.socketId = 0;
ScratchLink.sockets = new Map();

window.WebSocket = ScratchLink;

ViewController.swift

//
//  ViewController.swift
//  WebSocketInjection
//
//  Created by Shinichiro Oba on 2021/03/18.
//

import UIKit
import WebKit

class ViewController: UIViewController {
    
    @IBOutlet weak var webView: WKWebView!
    
    override func viewDidLoad() {
        super.viewDidLoad()
        
        let js = getInjectCode()
        let script = WKUserScript(source: js, injectionTime: .atDocumentEnd, forMainFrameOnly: false)
        webView.configuration.userContentController.addUserScript(script)
        webView.configuration.userContentController.add(self, name: "callbackHandler")
        
        let url = URL(string: "https://scratch.mit.edu/projects/editor/")!
        let request = URLRequest(url: url)
        webView.load(request)
    }
    
    func getInjectCode() -> String {
        guard let filepath = Bundle.main.path(forResource: "inject-scratch-link", ofType: "js") else { return "" }
        return (try? String(contentsOfFile: filepath)) ?? ""
    }
}

struct RPC: Codable {
    let method: Method
    let socketId: Int
    let url: String?
    let jsonrpc: JSONRPC?
    
    enum Method: String, Codable {
        case open
        case close
        case send
    }
    
    struct JSONRPC: Codable {
        let id: Int?
        let method: String
    }
}

extension ViewController: WKScriptMessageHandler {
    
    func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) {
        print(message.name, message.body)
        
        guard let jsonString = message.body as? String else { return }
        guard let jsonData = jsonString.data(using: .utf8) else { return }
        guard let json = try? JSONDecoder().decode(RPC.self, from: jsonData) else { return }
        
        let socketId = json.socketId
        
        if let id = json.jsonrpc?.id {
            sendJsonMessage(socketId: socketId, json: "{jsonrpc: \"2.0\", id: \(id), result: null}")
        }
        
        switch json.method {
        case .open:
            break
            
        case .close:
            break
            
        case .send:
            guard let method = json.jsonrpc?.method else { break }
            switch method {
            case "discover":
                DispatchQueue.main.asyncAfter(deadline: DispatchTime.now() + 1) { [weak self] in
                    self?.sendJsonMessage(socketId: socketId, json: "{jsonrpc: \"2.0\", method: \"didDiscoverPeripheral\", params: {name: \"Dummy Device 1\", peripheralId: \"11111111-AAAA-BBBB-CCCC-DDDDDDDDDDDD\", rssi: -10}}")
                }
                DispatchQueue.main.asyncAfter(deadline: DispatchTime.now() + 2) { [weak self] in
                    self?.sendJsonMessage(socketId: socketId, json: "{jsonrpc: \"2.0\", method: \"didDiscoverPeripheral\", params: {name: \"Dummy Device 2\", peripheralId: \"22222222-AAAA-BBBB-CCCC-DDDDDDDDDDDD\", rssi: -50}}")
                }
            default:
                break
            }
        }
    }
    
    private func sendJsonMessage(socketId: Int, json: String) {
        let js = "ScratchLink.sockets.get(\(socketId)).handleMessage(" + json + ")"
        webView.evaluateJavaScript(js)
    }
}