Merge pull request #238 from FlowFuse/dev-editor-rework

Save affinity cookie locally so it can be reused between restarts
FlowFuse · Feb 29, 2024 · 0b96755 · 0b96755
2 parents a7af207 + b62f339
commit 0b96755
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 15 deletions.
diff --git a/lib/agent.js b/lib/agent.js
@@ -32,6 +32,7 @@ class Agent {
         // that the first MQTT check-in will trigger a response
         this.currentState = 'unknown'
         this.editorToken = null
+        this.editorAffinity = null
         // ensure licensed property is present (default to null)
         if (utils.hasProperty(this.config, 'licensed') === false) {
             this.config.licensed = null
@@ -52,6 +53,7 @@ class Agent {
                     this.currentProject = null
                     this.currentApplication = null
                     this.editorToken = null
+                    this.editorAffinity = null
                 } else {
                     // New format
                     this.currentApplication = config.project ? null : (config.application || null)
@@ -61,6 +63,7 @@ class Agent {
                     this.currentMode = config.mode || 'autonomous'
                     this.config.licensed = config.licensed || null
                     this.editorToken = config.editorToken || null
+                    this.editorAffinity = config.editorAffinity || null
                 }
                 this.printAgentStatus()
             } catch (err) {
@@ -107,7 +110,8 @@ class Agent {
             settings: this.currentSettings,
             mode: this.currentMode,
             licensed: this.config.licensed,
-            editorToken: this.editorToken
+            editorToken: this.editorToken,
+            editorAffinity: this.editorAffinity
         }))
     }
 
@@ -167,7 +171,7 @@ class Agent {
         if (this.updating) {
             return null
         }
-        return {
+        const state = {
             ownerType: this.currentOwnerType,
             project: this.currentProject || null,
             application: this.currentApplication || null,
@@ -182,6 +186,10 @@ class Agent {
             agentVersion: this.config.version,
             licensed: this.config.licensed
         }
+        if (this.currentMode === 'developer' && this.editorToken && this.editorAffinity) {
+            state.affinity = this.editorAffinity
+        }
+        return state
     }
 
     /**
@@ -253,7 +261,8 @@ class Agent {
                     await this.saveProject()
                 } else {
                     // exiting developer mode
-                    this.editorToken = null // clear the discarded token
+                    this.editorToken = null
+                    this.editorAffinity = null
                     let _launcher = this.launcher
                     if (!_launcher) {
                         // create a temporary launcher to read the current snapshot on disk
@@ -366,6 +375,7 @@ class Agent {
                 this.currentSettings = null
                 this.currentMode = null
                 this.editorToken = null
+                this.editorAffinity = null
                 await this.saveProject()
                 this.currentState = 'stopped'
                 this.updating = false
@@ -535,7 +545,7 @@ class Agent {
                         this.mqttClient.setProject(this.currentProject)
                         this.mqttClient.setApplication(this.currentApplication)
                         if (developerMode && this.editorToken) {
-                            this.mqttClient.startTunnel(this.editorToken)
+                            this.mqttClient.startTunnel(this.editorToken, this.editorAffinity)
                         }
                     }
                     this.checkIn(2)
@@ -583,7 +593,7 @@ class Agent {
                             this.mqttClient.setProject(this.currentProject)
                             this.mqttClient.setApplication(this.currentApplication)
                             if (developerMode && this.editorToken) {
-                                this.mqttClient.startTunnel(this.editorToken)
+                                this.mqttClient.startTunnel(this.editorToken, this.editorAffinity)
                             }
                         }
                         this.checkIn(2)
@@ -689,9 +699,10 @@ class Agent {
         })
     }
 
-    async saveEditorToken (token) {
-        const changed = this.editorToken !== token
+    async saveEditorToken (token, affinity) {
+        const changed = (this.editorToken !== token || this.editorAffinity !== affinity)
         this.editorToken = token
+        this.editorAffinity = affinity
         if (changed) {
             await this.saveProject()
         }

diff --git a/lib/editor/tunnel.js b/lib/editor/tunnel.js
@@ -22,7 +22,7 @@ class EditorTunnel {
         this.port = config.port
         this.config = config
         this.options = options || {}
-        this.affinity = undefined
+        this.affinity = this.options.affinity
 
         // How long to wait before attempting to reconnect. Start at 500ms - back
         // off if connect fails

diff --git a/lib/mqtt.js b/lib/mqtt.js
@@ -100,9 +100,11 @@ class MQTTClient {
                     this.logEnabled = false
                     return
                 } else if (msg.command === 'startEditor') {
-                    await this.startTunnel(msg.payload?.token, msg)
+                    await this.startTunnel(msg.payload?.token, this.agent.editorAffinity || null, msg)
                     return
                 } else if (msg.command === 'stopEditor') {
+                    // Clear the saved token
+                    await this.saveEditorToken(null, null)
                     if (this.tunnel) {
                         info('Disabling remote editor access')
                         this.tunnel.close()
@@ -266,7 +268,7 @@ class MQTTClient {
         })
     }
 
-    async startTunnel (token, msg) {
+    async startTunnel (token, affinity, msg) {
         info('Enabling remote editor access')
         try {
             if (this.tunnel) {
@@ -282,17 +284,23 @@ class MQTTClient {
             }
 
             // * Enable Device Editor (Step 6) - (forge:MQTT->device) Create the tunnel on the device
-            this.tunnel = EditorTunnel.create(this.config, { token })
+            this.tunnel = EditorTunnel.create(this.config, { token, affinity })
 
             // * Enable Device Editor (Step 7) - (device) Begin the device tunnel connect process
             const result = await this.tunnel.connect()
 
             // store the token for later use (i.e. device agent is restarted)
-            await this.saveEditorToken(result ? token : null)
+            if (result) {
+                await this.saveEditorToken(token, this.tunnel.affinity)
+            } else {
+                // Failed to connect - clear the token/affinity so it can be
+                // refreshed
+                await this.saveEditorToken(null, null)
+            }
 
             if (msg) {
                 // * Enable Device Editor (Step 10) - (device->forge:MQTT) Send a response to the platform
-                this.sendCommandResponse(msg, { connected: result, token })
+                this.sendCommandResponse(msg, { connected: result, token, affinity: this.tunnel.affinity })
             }
         } catch (err) {
             warn(`Error starting editor tunnel: ${err}`)
@@ -303,8 +311,8 @@ class MQTTClient {
         this.sendStatus()
     }
 
-    async saveEditorToken (token) {
-        await this.agent?.saveEditorToken(token)
+    async saveEditorToken (token, affinity) {
+        await this.agent?.saveEditorToken(token, affinity)
     }
 }