peardock/server/rpc/binary-stream.js

/**
 * Dedicated binary transfer over protomux-rpc chunked events with backpressure.
 *
 * Protocol (control plane stays JSON RPC; payload is base64 chunks on push):
 *   Client → binaryStreamOpen  { kind, id?, maxBytes?, direction: 'download'|'upload' }
 *   Server → push:binaryChunk  { streamId, index, data, done, totalBytes? }
 *   Client → binaryStreamChunk { streamId, index, data }  (upload)
 *   Client → binaryStreamClose { streamId }
 *
 * This is the production binary path; hyperschema raw encodings can replace base64 later
 * without changing method names.
 */
import { docker } from '../services/docker.js'
import { Pushes, Methods } from '../../shared/protocol.js'
import logger from '../utils/logger.js'

const DEFAULT_CHUNK = 256 * 1024
const DEFAULT_MAX = 100 * 1024 * 1024

/**
 * @param {import('./session.js').PeerSession} session
 */
export function registerBinaryStreamHandlers(session) {
  /** @type {Map<string, { kind: string, buffers: Buffer[], total: number, maxBytes: number, meta: object }>} */
  const uploads = new Map()

  session.respond(Methods.binaryStreamOpen || 'binaryStreamOpen', async (args) => {
    const kind = args.kind
    const direction = args.direction || (kind === 'imageLoad' || kind === 'upload' ? 'upload' : 'download')
    const maxBytes = Math.min(Number(args.maxBytes) || DEFAULT_MAX, 200 * 1024 * 1024)
    const chunkSize = Math.min(Number(args.chunkSize) || DEFAULT_CHUNK, 512 * 1024)
    const streamId = `bs-${Date.now()}-${Math.random().toString(36).slice(2, 10)}`

    if (direction === 'upload') {
      uploads.set(streamId, { kind, buffers: [], total: 0, maxBytes, meta: args })
      session.state.set(`bstream:${streamId}`, { kind: 'upload', streamId })
      return { success: true, streamId, direction: 'upload', maxBytes, chunkSize }
    }

    // Download path: open engine stream and push chunks
    let readable
    let meta = {}
    if (kind === 'imageSave') {
      if (!args.id) throw new Error('id required for imageSave')
      readable = await docker.getImage(args.id).get()
      meta = { imageId: args.id }
    } else if (kind === 'containerExport') {
      if (!args.id) throw new Error('id required for containerExport')
      readable = await docker.getContainer(args.id).export()
      meta = { containerId: args.id }
    } else {
      throw new Error(`Unsupported download kind: ${kind}`)
    }

    // Async push — don't block response longer than first bytes
    setImmediate(() => {
      pumpDownload(session, streamId, readable, { chunkSize, maxBytes, kind, ...meta }).catch((err) => {
        logger.error('binary stream download failed', { streamId, error: err.message })
        try {
          session.push(Pushes.binaryChunk, {
            type: 'binaryChunk',
            streamId,
            kind,
            error: err.message,
            done: true,
            index: -1,
          })
        } catch {
          // ignore
        }
      })
    })

    return { success: true, streamId, direction: 'download', maxBytes, chunkSize, kind }
  })

  session.respond(Methods.binaryStreamChunk || 'binaryStreamChunk', async (args) => {
    const streamId = args.streamId
    const entry = uploads.get(streamId)
    if (!entry) throw new Error('Unknown upload streamId')
    if (!args.data) throw new Error('data required')
    const buf = Buffer.from(args.data, args.encoding === 'utf8' ? 'utf8' : 'base64')
    entry.total += buf.length
    if (entry.total > entry.maxBytes) {
      uploads.delete(streamId)
      throw new Error(`Upload exceeds maxBytes (${entry.maxBytes})`)
    }
    entry.buffers.push(buf)
    return { success: true, streamId, received: entry.total, index: args.index }
  })

  session.respond(Methods.binaryStreamClose || 'binaryStreamClose', async (args) => {
    const streamId = args.streamId
    const entry = uploads.get(streamId)
    if (!entry) {
      // download close is ack-only
      return { success: true, streamId, closed: true }
    }
    uploads.delete(streamId)
    const buf = Buffer.concat(entry.buffers)
    if (entry.kind === 'imageLoad' || entry.kind === 'upload') {
      if (!buf.length) throw new Error('No data received')
      const stream = await docker.loadImage(buf)
      await new Promise((resolve, reject) => {
        docker.modem.followProgress(stream, (err) => (err ? reject(err) : resolve()))
      })
      return {
        success: true,
        streamId,
        kind: entry.kind,
        size: buf.length,
        message: 'Image loaded via binary stream',
      }
    }
    return { success: true, streamId, size: buf.length, message: 'Upload closed' }
  })
}

/**
 * @param {import('./session.js').PeerSession} session
 * @param {string} streamId
 * @param {import('stream').Readable} readable
 * @param {{ chunkSize: number, maxBytes: number, kind: string }} opts
 */
async function pumpDownload(session, streamId, readable, opts) {
  let total = 0
  let index = 0
  let pending = Buffer.alloc(0)

  await new Promise((resolve, reject) => {
    readable.on('data', (chunk) => {
      total += chunk.length
      if (total > opts.maxBytes) {
        readable.destroy()
        reject(new Error(`Download exceeds maxBytes (${opts.maxBytes})`))
        return
      }
      pending = Buffer.concat([pending, chunk])
      while (pending.length >= opts.chunkSize) {
        const slice = pending.subarray(0, opts.chunkSize)
        pending = pending.subarray(opts.chunkSize)
        session.push(Pushes.binaryChunk, {
          type: 'binaryChunk',
          streamId,
          kind: opts.kind,
          index,
          encoding: 'base64',
          data: slice.toString('base64'),
          done: false,
          totalBytes: null,
        })
        index += 1
      }
    })
    readable.on('end', () => {
      if (pending.length) {
        session.push(Pushes.binaryChunk, {
          type: 'binaryChunk',
          streamId,
          kind: opts.kind,
          index,
          encoding: 'base64',
          data: pending.toString('base64'),
          done: false,
          totalBytes: total,
        })
        index += 1
      }
      session.push(Pushes.binaryChunk, {
        type: 'binaryChunk',
        streamId,
        kind: opts.kind,
        index,
        encoding: 'base64',
        data: '',
        done: true,
        totalBytes: total,
      })
      resolve()
    })
    readable.on('error', reject)
  })
}