RPC work

DynamicIsland/services/Extensions/ExtensionRPCServer.swift

@@ -0,0 +1,294 @@
+/*
+ * Atoll (DynamicIsland)
+ * Copyright (C) 2024-2026 Atoll Contributors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see <https://www.gnu.org/licenses/>.
+ */
+
+import Foundation
+import Network
+import Defaults
+
+/// WebSocket server for Atoll RPC.
+/// Uses Apple's Network.framework (`NWListener`) — no external dependencies.
+/// Listens on localhost:9020 for JSON-RPC 2.0 requests over WebSocket.
+@MainActor
+final class ExtensionRPCServer {
+    static let shared = ExtensionRPCServer()
+
+    private var listener: NWListener?
+    private var connections: [UUID: RPCClientConnection] = [:]
+    private let port: UInt16 = 9020
+    private let queue = DispatchQueue(label: "com.ebullioscopic.Atoll.rpc.server", qos: .userInitiated)
+    private let decoder = JSONDecoder()
+    private let encoder = JSONEncoder()
+
+    private init() {}
+
+    // MARK: - Lifecycle
+
+    func start() {
+        guard listener == nil else {
+            logDiagnostics("RPC server already running")
+            return
+        }
+
+        let params = NWParameters(tls: nil)
+        let wsOptions = NWProtocolWebSocket.Options()
+        wsOptions.autoReplyPing = true
+        params.defaultProtocolStack.applicationProtocols.insert(wsOptions, at: 0)
+
+        do {
+            listener = try NWListener(using: params, on: NWEndpoint.Port(integerLiteral: port))
+        } catch {
+            Logger.log("Failed to create RPC listener: \(error.localizedDescription)", category: .extensions)
+            return
+        }
+
+        listener?.stateUpdateHandler = { [weak self] state in
+            Task { @MainActor in
+                self?.handleListenerState(state)
+            }
+        }
+
+        listener?.newConnectionHandler = { [weak self] connection in
+            Task { @MainActor in
+                self?.handleNewConnection(connection)
+            }
+        }
+
+        listener?.start(queue: queue)
+    }
+
+    func stop() {
+        listener?.cancel()
+        listener = nil
+        for (_, conn) in connections {
+            conn.connection.cancel()
+        }
+        connections.removeAll()
+        logDiagnostics("RPC server stopped")
+    }
+
+    // MARK: - Client Notifications
+
+    func notifyActivityDismiss(bundleIdentifier: String, activityID: String) {
+        sendNotification(
+            to: bundleIdentifier,
+            method: "atoll.activityDidDismiss",
+            params: [
+                "bundleIdentifier": .string(bundleIdentifier),
+                "activityID": .string(activityID)
+            ]
+        )
+    }
+
+    func notifyWidgetDismiss(bundleIdentifier: String, widgetID: String) {
+        sendNotification(
+            to: bundleIdentifier,
+            method: "atoll.widgetDidDismiss",
+            params: [
+                "bundleIdentifier": .string(bundleIdentifier),
+                "widgetID": .string(widgetID)
+            ]
+        )
+    }
+
+    func notifyNotchExperienceDismiss(bundleIdentifier: String, experienceID: String) {
+        sendNotification(
+            to: bundleIdentifier,
+            method: "atoll.notchExperienceDidDismiss",
+            params: [
+                "bundleIdentifier": .string(bundleIdentifier),
+                "experienceID": .string(experienceID)
+            ]
+        )
+    }
+
+    func notifyAuthorizationChange(bundleIdentifier: String, isAuthorized: Bool) {
+        sendNotification(
+            to: bundleIdentifier,
+            method: "atoll.authorizationDidChange",
+            params: [
+                "bundleIdentifier": .string(bundleIdentifier),
+                "isAuthorized": .bool(isAuthorized)
+            ]
+        )
+    }
+
+    // MARK: - Connection Handling
+
+    private func handleListenerState(_ state: NWListener.State) {
+        switch state {
+        case .ready:
+            Logger.log("Started Atoll RPC WebSocket server on port \(port)", category: .extensions)
+        case .failed(let error):
+            Logger.log("RPC server failed: \(error.localizedDescription)", category: .extensions)
+            // Attempt restart after delay
+            DispatchQueue.main.asyncAfter(deadline: .now() + 3) { [weak self] in
+                self?.listener = nil
+                self?.start()
+            }
+        case .cancelled:
+            logDiagnostics("RPC server listener cancelled")
+        default:
+            break
+        }
+    }
+
+    private func handleNewConnection(_ nwConnection: NWConnection) {
+        let connID = UUID()
+        let clientConn = RPCClientConnection(
+            id: connID,
+            connection: nwConnection,
+            bundleIdentifier: nil
+        )
+        connections[connID] = clientConn
+
+        nwConnection.stateUpdateHandler = { [weak self] state in
+            Task { @MainActor in
+                self?.handleConnectionState(connID: connID, state: state)
+            }
+        }
+
+        nwConnection.start(queue: queue)
+        receiveMessage(connID: connID)
+        logDiagnostics("RPC client connected (id: \(connID.uuidString.prefix(8)))")
+    }
+
+    private func handleConnectionState(connID: UUID, state: NWConnection.State) {
+        switch state {
+        case .failed, .cancelled:
+            connections.removeValue(forKey: connID)
+            logDiagnostics("RPC client disconnected (id: \(connID.uuidString.prefix(8)))")
+        default:
+            break
+        }
+    }
+
+    private func receiveMessage(connID: UUID) {
+        guard let clientConn = connections[connID] else { return }
+        let connection = clientConn.connection
+
+        connection.receiveMessage { [weak self] content, context, isComplete, error in
+            guard let self else { return }
+
+            if let error {
+                Task { @MainActor in
+                    self.logDiagnostics("RPC receive error for \(connID.uuidString.prefix(8)): \(error.localizedDescription)")
+                    self.connections.removeValue(forKey: connID)
+                }
+                return
+            }
+
+            if let data = content, !data.isEmpty {
+                Task { @MainActor in
+                    self.processMessage(data: data, connID: connID)
+                }
+            }
+
+            // Continue receiving
+            Task { @MainActor in
+                self.receiveMessage(connID: connID)
+            }
+        }
+    }
+
+    private func processMessage(data: Data, connID: UUID) {
+        guard var clientConn = connections[connID] else { return }
+
+        // Parse JSON-RPC request
+        guard let request = try? decoder.decode(RPCRequest.self, from: data) else {
+            let errorResponse = RPCErrorResponse(
+                error: RPCErrorObject(code: RPCErrorCode.parseError, message: "Invalid JSON-RPC request"),
+                id: nil
+            )
+            sendResponse(errorResponse, to: connID)
+            return
+        }
+
+        // Resolve bundle identifier from first authorization request
+        if clientConn.bundleIdentifier == nil,
+           let params = request.params,
+           let bi = params["bundleIdentifier"]?.stringValue {
+            clientConn.bundleIdentifier = bi
+            connections[connID] = clientConn
+        }
+
+        let service = ExtensionRPCService(
+            bundleIdentifier: clientConn.bundleIdentifier ?? "unknown",
+            server: self
+        )
+
+        let responseData = service.handleRequest(request)
+        sendRawData(responseData, to: connID)
+    }
+
+    // MARK: - Send Helpers
+
+    private func sendResponse(_ response: Codable, to connID: UUID) {
+        guard let data = try? encoder.encode(response) else { return }
+        sendRawData(data, to: connID)
+    }
+
+    private func sendRawData(_ data: Data, to connID: UUID) {
+        guard let clientConn = connections[connID] else { return }
+
+        let metadata = NWProtocolWebSocket.Metadata(opcode: .text)
+        let context = NWConnection.ContentContext(identifier: "rpc-response", metadata: [metadata])
+
+        clientConn.connection.send(
+            content: data,
+            contentContext: context,
+            isComplete: true,
+            completion: .contentProcessed { error in
+                if let error {
+                    Task { @MainActor in
+                        self.logDiagnostics("RPC send error: \(error.localizedDescription)")
+                    }
+                }
+            }
+        )
+    }
+
+    private func sendNotification(to bundleIdentifier: String, method: String, params: [String: RPCValue]) {
+        let notification = RPCNotification(method: method, params: params)
+        guard let data = try? encoder.encode(notification) else { return }
+
+        for (_, clientConn) in connections where clientConn.bundleIdentifier == bundleIdentifier {
+            let metadata = NWProtocolWebSocket.Metadata(opcode: .text)
+            let context = NWConnection.ContentContext(identifier: "rpc-notification", metadata: [metadata])
+
+            clientConn.connection.send(
+                content: data,
+                contentContext: context,
+                isComplete: true,
+                completion: .contentProcessed { _ in }
+            )
+        }
+    }
+
+    private func logDiagnostics(_ message: String) {
+        guard Defaults[.extensionDiagnosticsLoggingEnabled] else { return }
+        Logger.log(message, category: .extensions)
+    }
+}
+
+// MARK: - Client Connection
+
+struct RPCClientConnection {
+    let id: UUID
+    let connection: NWConnection
+    var bundleIdentifier: String?
+}