@adobe/aio-lib-sandbox 0.1.0-alpha.4 → 0.1.0-alpha.5

package/src/ws.js

@@ -12,6 +12,7 @@ governing permissions and limitations under the License.
 const WebSocket = require('ws')
 const {
   SandboxClientError,
+  SandboxCommandNotFoundError,
   SandboxUnauthorizedError,
   SandboxWebSocketError
 } = require('./errors')
@@ -36,11 +37,27 @@ class SandboxSocket {
 
     this.socket = null
     this.connectPromise = null
-
-    /** @type {Map<string, {resolve: Function, reject: Function, stdout: string, stderr: string, onOutput: Function|undefined, timeout: any}>} */
+    this.intentionalClose = false
+
+    /**
+     * Pending exec entries
+     *
+     * @type {Map<string, {
+     *   resolve: Function, reject: Function,
+     *   waitResolve: Function|null, waitReject: Function|null,
+     *   stdout: string, stderr: string,
+     *   onOutput: Function|undefined, timeout: any,
+     *   detached: boolean, resolved: boolean
+     * }>}
+     */
     this.pendingExecs = new Map()
     /** @type {Map<string, {resolve: Function, reject: Function}>} */
     this.pendingFileOps = new Map()
+    /**
+     * Pending exec.get operations keyed by execId.
+     * @type {Map<string, {resolve: Function, reject: Function, onOutput: Function|undefined, sandbox: object}>}
+     */
+    this.pendingGetOps = new Map()
   }
 
   /**
@@ -84,6 +101,10 @@ class SandboxSocket {
       const onClose = (code) => {
         cleanup()
         this.connectPromise = null
+        if (this.intentionalClose) {
+          resolve()
+          return
+        }
         reject(this.createCloseError(code))
       }
 
@@ -129,10 +150,28 @@ class SandboxSocket {
     this.socket.send(JSON.stringify(frame))
   }
 
+  /**
+   * Marks the next socket close as expected by the caller.
+   */
+  beginIntentionalClose () {
+    this.intentionalClose = true
+  }
+
+  /**
+   * Clears a previously requested intentional close.
+   */
+  cancelIntentionalClose () {
+    this.intentionalClose = false
+  }
+
   /**
    * Closes the underlying socket.
+   *
+   * @param {object} [options]
+   * @param {boolean} [options.intentional] whether pending work should be drained without error
    */
-  close () {
+  close ({ intentional = false } = {}) {
+    if (intentional) this.beginIntentionalClose()
     this.socket?.close()
   }
 
@@ -140,6 +179,60 @@ class SandboxSocket {
   // Pending operation helpers
   // ------------------------------------------------------------------
 
+  /**
+   * Registers a pending exec, sends the frame, and returns the promises 
+   * that will be settled by subsequent result or ack frames (for detached)
+   *
+   * @param {string} execId
+   * @param {object} frame the frame to send (must include `type` and `execId`)
+   * @param {{ detached: boolean, onOutput?: Function }} options
+   * @returns {{ ackPromise: Promise, waitPromise: Promise|null }}
+   */
+  sendExec (execId, frame, { detached, onOutput }) {
+    let waitResolve, waitReject, waitPromise
+    if (detached) {
+      waitPromise = new Promise((res, rej) => { waitResolve = res; waitReject = rej })
+    }
+
+    let resolve, reject
+    const ackPromise = new Promise((res, rej) => { resolve = res; reject = rej })
+
+    this.pendingExecs.set(execId, {
+      resolve,
+      reject,
+      waitResolve: waitResolve || null,
+      waitReject: waitReject || null,
+      _waitPromise: waitPromise || null,
+      stdout: '',
+      stderr: '',
+      onOutput: onOutput || null,
+      timeout: undefined,
+      detached,
+      resolved: false
+    })
+
+    try {
+      this.send(frame)
+    } catch (error) {
+      this.rejectExec(execId, new SandboxWebSocketError(
+        `Could not send exec frame: ${error.message}`
+      ))
+    }
+
+    return { ackPromise, waitPromise: waitPromise || null }
+  }
+
+  /**
+   * Stores a timer handle on a pending exec entry so it can be cleared on completion.
+   *
+   * @param {string} execId
+   * @param {ReturnType<setTimeout>} handle
+   */
+  setExecTimeout (execId, handle) {
+    const entry = this.pendingExecs.get(execId)
+    if (entry) entry.timeout = handle
+  }
+
   /**
    * Rejects and removes a pending exec, clearing its timeout.
    *
@@ -151,7 +244,13 @@ class SandboxSocket {
     if (!pending) return
     this.pendingExecs.delete(execId)
     clearTimeout(pending.timeout)
-    pending.reject(error)
+
+    // For detached execs, the first promise is always resolved, so reject the wait promise instead
+    if (pending.detached && pending.resolved) {
+      if (pending.waitReject) pending.waitReject(error)
+    } else {
+      pending.reject(error)
+    }
   }
 
   /**
@@ -167,10 +266,45 @@ class SandboxSocket {
     pending.reject(error)
   }
 
+  /**
+   * Resolves and removes a pending exec during an intentional sandbox shutdown.
+   *
+   * @param {string} execId
+   */
+  resolveExecOnIntentionalClose (execId) {
+    const pending = this.pendingExecs.get(execId)
+    if (!pending) return
+    this.pendingExecs.delete(execId)
+    clearTimeout(pending.timeout)
+
+    const result = { exitCode: null, destroyed: true }
+    if (pending.detached) {
+      if (!pending.resolved) {
+        pending.resolved = true
+        pending.resolve({ pid: undefined, startedAt: undefined, destroyed: true })
+      }
+      if (pending.waitResolve) pending.waitResolve(result)
+      return
+    }
+
+    pending.resolve({
+      execId,
+      stdout: pending.stdout,
+      stderr: pending.stderr,
+      ...result
+    })
+  }
+
   handleMessage (message) {
     const frame = this.parseFrame(message)
     if (!frame || this.isAuthAckFrame(frame)) return
 
+    if (frame.type === 'exec.info' ||
+        (frame.type === 'error' && this.pendingGetOps.has(frame.execId))) {
+      this.handleGetFrame(frame)
+      return
+    }
+
     if (this.pendingFileOps.has(frame.execId)) {
       this.handleFileFrame(frame)
       return
@@ -182,6 +316,24 @@ class SandboxSocket {
   }
 
   handleClose (code) {
+    if (this.intentionalClose) {
+      for (const execId of [...this.pendingExecs.keys()]) {
+        this.resolveExecOnIntentionalClose(execId)
+      }
+      for (const [, pending] of [...this.pendingFileOps.entries()]) {
+        pending.resolve(undefined)
+      }
+      for (const [, pending] of [...this.pendingGetOps.entries()]) {
+        pending.resolve(null)
+      }
+      this.pendingFileOps.clear()
+      this.pendingGetOps.clear()
+      this.connectPromise = null
+      this.socket = null
+      this.intentionalClose = false
+      return
+    }
+
     const error = this.createCloseError(code)
     for (const execId of [...this.pendingExecs.keys()]) {
       this.rejectExec(execId, error)
@@ -189,6 +341,10 @@ class SandboxSocket {
     for (const execId of [...this.pendingFileOps.keys()]) {
       this.rejectFileOp(execId, error)
     }
+    for (const [, pending] of [...this.pendingGetOps.entries()]) {
+      pending.reject(error)
+    }
+    this.pendingGetOps.clear()
     this.connectPromise = null
     this.socket = null
   }
@@ -236,15 +392,32 @@ class SandboxSocket {
       return
     }
 
+    // For detached cmds, we will resolve with the server response, which is all the command info needed for a
+    // handle. The entry stays in pendingExecs to receive subsequent output and exec.exit.
+    if (frame.type === 'exec.detached') {
+      clearTimeout(pending.timeout)
+      pending.timeout = undefined
+      pending.resolved = true
+      pending.resolve({ pid: frame.pid, startedAt: frame.startedAt })
+      return
+    }
+
     if (frame.type === 'exec.exit') {
       this.pendingExecs.delete(frame.execId)
       clearTimeout(pending.timeout)
-      pending.resolve({
-        execId: frame.execId,
-        stdout: pending.stdout,
-        stderr: pending.stderr,
-        exitCode: frame.exitCode
-      })
+      if (pending.detached && pending.resolved) {
+        // For detached, the initial ack promise already resolved, so we resolve the wait promise
+        if (pending.waitResolve) {
+          pending.waitResolve({ exitCode: frame.exitCode })
+        }
+      } else {
+        pending.resolve({
+          execId: frame.execId,
+          stdout: pending.stdout,
+          stderr: pending.stderr,
+          exitCode: frame.exitCode
+        })
+      }
       return
     }
 
@@ -292,6 +465,133 @@ class SandboxSocket {
       ))
     }
   }
+
+  /**
+   * Handles exec.info (response to exec.get) and error frames routed from handleMessage.
+   *
+   * @param {object} frame
+   */
+  handleGetFrame (frame) {
+    const pending = this.pendingGetOps.get(frame.execId)
+    if (!pending) return
+
+    if (frame.type === 'exec.info') {
+      this.resolveGetOp(frame, pending)
+      return
+    }
+
+    if (frame.type === 'error') {
+      this.rejectGetOp(frame, pending)
+    }
+  }
+
+  /**
+   * Resolves a pending exec.get by building a command handle and resolving the caller's promise.
+   *
+   * @param {object} frame exec.info frame
+   * @param {object} pending entry from pendingGetOps
+   */
+  resolveGetOp (frame, pending) {
+    this.pendingGetOps.delete(frame.execId)
+    const waitPromise = this.resolveExecEntry(frame, pending)
+    const commandObj = this.buildCommandObject(frame, waitPromise, pending.sandbox)
+    pending.resolve(commandObj)
+  }
+
+  /**
+   * Rejects a pending exec.get with a not-found error.
+   *
+   * @param {object} frame error frame
+   * @param {object} pending entry from pendingGetOps
+   */
+  rejectGetOp (frame, pending) {
+    this.pendingGetOps.delete(frame.execId)
+    pending.reject(new SandboxCommandNotFoundError(
+      frame.message || `No running process for execId '${frame.execId}'`
+    ))
+  }
+
+  /**
+   * Returns the wait promise for the exec, either by reusing an existing pendingExecs entry
+   * (same session) or by registering a fresh reattached entry (new session / previous connection).
+   *
+   * @param {object} frame exec.info frame
+   * @param {object} pending entry from pendingGetOps
+   * @returns {Promise}
+   */
+  resolveExecEntry (frame, pending) {
+    const existingExec = this.pendingExecs.get(frame.execId)
+    if (existingExec) {
+      this.mergeOnOutputCallback(existingExec, pending.onOutput)
+      return existingExec._waitPromise
+    }
+    return this.registerReattachedExec(frame, pending.onOutput)
+  }
+
+  /**
+   * Appends `onOutput` to an existing exec entry's callback chain, preserving the previous handler.
+   *
+   * @param {object} existingExec entry from pendingExecs
+   * @param {Function|undefined} onOutput new callback to add
+   */
+  mergeOnOutputCallback (existingExec, onOutput) {
+    if (!onOutput) return
+    const prev = existingExec.onOutput
+    existingExec.onOutput = (data, stream) => {
+      if (prev) prev(data, stream)
+      onOutput(data, stream)
+    }
+  }
+
+  /**
+   * Creates a fresh pendingExecs entry for a process reattached from a previous connection,
+   * registers it, and returns its wait promise.
+   *
+   * @param {object} frame exec.info frame
+   * @param {Function|undefined} onOutput output callback
+   * @returns {Promise}
+   */
+  registerReattachedExec (frame, onOutput) {
+    let waitResolve, waitReject
+    const waitPromise = new Promise((res, rej) => { waitResolve = res; waitReject = rej })
+    this.pendingExecs.set(frame.execId, {
+      resolve: () => {},
+      reject: () => {},
+      waitResolve,
+      waitReject,
+      _waitPromise: waitPromise,
+      stdout: '',
+      stderr: '',
+      onOutput: onOutput || null,
+      timeout: undefined,
+      detached: frame.detached,
+      resolved: true
+    })
+    return waitPromise
+  }
+
+  /**
+   * Builds the command handle object returned to the caller of exec.get.
+   *
+   * @param {object} frame exec.info frame
+   * @param {Promise} waitPromise resolves when the process exits
+   * @param {object} sandbox Sandbox instance for delegating control operations
+   * @returns {object}
+   */
+  buildCommandObject (frame, waitPromise, sandbox) {
+    const { execId } = frame
+    return {
+      execId,
+      command: frame.command,
+      pid: frame.pid,
+      startedAt: frame.startedAt,
+      detached: frame.detached,
+      wait: () => waitPromise,
+      writeStdin: (data) => sandbox.writeStdin(execId, data),
+      closeStdin: () => sandbox.closeStdin(execId),
+      kill: (signal) => sandbox.kill(execId, signal)
+    }
+  }
 }
 
 module.exports = { SandboxSocket }