From 6892b443e0b527d8f2725c6259fa954bf4553505 Mon Sep 17 00:00:00 2001
From: Koushik Dutta <koushd@gmail.com>
Date: Tue, 25 Jun 2024 00:16:20 -0700
Subject: [PATCH] server: fix plugin fork storage desync

---
 server/src/plugin/plugin-remote-worker.ts | 21 ++++++++++++++++++---
 server/src/plugin/plugin-remote.ts        |  4 ++--
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/server/src/plugin/plugin-remote-worker.ts b/server/src/plugin/plugin-remote-worker.ts
index 68df01cc93..2bd826ce30 100644
--- a/server/src/plugin/plugin-remote-worker.ts
+++ b/server/src/plugin/plugin-remote-worker.ts
@@ -55,6 +55,8 @@ export function startPluginRemote(mainFilename: string, pluginId: string, peerSe
 
     let postInstallSourceMapSupport: (scrypted: ScryptedStatic) => void;
 
+    const forks = new Set<PluginRemote>();
+
     attachPluginRemote(peer, {
         createMediaManager: async (sm, dm) => {
             systemManager = sm;
@@ -62,8 +64,21 @@ export function startPluginRemote(mainFilename: string, pluginId: string, peerSe
             return new MediaManagerImpl(systemManager, dm);
         },
         onGetRemote: async (_api, _pluginId) => {
-            api = _api;
+            class PluginForkableAPI extends PluginAPIProxy {
+                [RpcPeer.PROPERTY_PROXY_ONEWAY_METHODS] = (_api as any)[RpcPeer.PROPERTY_PROXY_ONEWAY_METHODS];
+
+                setStorage(nativeId: string, storage: { [key: string]: any; }): Promise<void> {
+                    const id = deviceManager.nativeIds.get(nativeId).id;
+                    for (const r of forks) {
+                        r.setNativeId(nativeId, id, storage);
+                    }
+                    return super.setStorage(nativeId, storage);
+                }
+            }
+
+            api = new PluginForkableAPI(_api);
             peer.selfName = pluginId;
+            return api;
         },
         getPluginConsole,
         getDeviceConsole,
@@ -298,7 +313,7 @@ export function startPluginRemote(mainFilename: string, pluginId: string, peerSe
                 process.send(options, socket);
             }
 
-            const forks = new Set<PluginRemote>();
+            const pluginRemoteAPI: PluginRemote = scrypted.pluginRemoteAPI;
 
             scrypted.fork = () => {
                 const ntw = new NodeThreadWorker(mainFilename, pluginId, {
@@ -322,7 +337,7 @@ export function startPluginRemote(mainFilename: string, pluginId: string, peerSe
 
                         setStorage(nativeId: string, storage: { [key: string]: any; }): Promise<void> {
                             const id = deviceManager.nativeIds.get(nativeId).id;
-                            (scrypted.pluginRemoteAPI as PluginRemote).setNativeId(nativeId, id, storage);
+                            pluginRemoteAPI.setNativeId(nativeId, id, storage);
                             for (const r of forks) {
                                 if (r === remote)
                                     continue;
diff --git a/server/src/plugin/plugin-remote.ts b/server/src/plugin/plugin-remote.ts
index cce73eb434..2f35492857 100644
--- a/server/src/plugin/plugin-remote.ts
+++ b/server/src/plugin/plugin-remote.ts
@@ -462,7 +462,7 @@ export interface PluginRemoteAttachOptions {
     getPluginConsole?: () => Console;
     getMixinConsole?: (id: string, nativeId?: ScryptedNativeId) => Console;
     onLoadZip?: (scrypted: ScryptedStatic, params: any, packageJson: any, getZip: () => Promise<Buffer>, zipOptions: PluginRemoteLoadZipOptions) => Promise<any>;
-    onGetRemote?: (api: PluginAPI, pluginId: string) => Promise<void>;
+    onGetRemote?: (api: PluginAPI, pluginId: string) => Promise<PluginAPI>;
 }
 
 export function attachPluginRemote(peer: RpcPeer, options?: PluginRemoteAttachOptions): Promise<ScryptedStatic> {
@@ -496,7 +496,7 @@ export function attachPluginRemote(peer: RpcPeer, options?: PluginRemoteAttachOp
             }
         });
 
-        await options?.onGetRemote?.(api, pluginId);
+        api = await options?.onGetRemote?.(api, pluginId) || api;
 
         const systemManager = new SystemManagerImpl();
         const deviceManager = new DeviceManagerImpl(systemManager, getDeviceConsole, getMixinConsole);