Patching Runno WASI to allow communications in raw byte streams

patch-wasi.js

import * as Runno from '@runno/wasi'
import wrapPollOneoff from './wrap-poll-oneoff'

const { Result } = Runno.WASISnapshotPreview1

/** @typedef {{ lenTotal: number, ptrlens: [number, number][] }} IovsDesc */

/**
 * @param {DataView} view
 * @param {number} iovs_ptr
 * @param {number} iovs_len
 * @returns {IovsDesc}
 */
function collectIOVectors(view, iovs_ptr, iovs_len) {
  /** @type {[number, number][]} */
  const ptrlens = []
  let lenTotal = 0

  for (let i = 0; i < iovs_len; i++) {
    const bufferPtr = view.getUint32(iovs_ptr, true)
    iovs_ptr += 4

    const bufferLen = view.getUint32(iovs_ptr, true)
    iovs_ptr += 4
    lenTotal += bufferLen

    ptrlens.push([bufferPtr, bufferLen])
  }

  return { lenTotal, ptrlens }
}

/**
 * @param {DataView} view
 * @param {IovsDesc} iovsDesc
 * @returns {Uint8Array}
 */
function readIOVectorsMerged(view, iovsDesc) {
  const source = new Uint8Array(view.buffer, view.byteOffset, view.byteLength)

  const result = new Uint8Array(iovsDesc.lenTotal)
  let written = 0
  for (const [ptr, len] of iovsDesc.ptrlens) {
    // XXX: is there a cleaner way?
    result.set(source.subarray(ptr, ptr + len), written)
    written += len
  }

  return result
}

/**
 * @param {Uint8Array} buf
 * @param {IovsDesc} iovsDesc
 * @param {Uint8Array} input
 */
function writeIntoIOVectors(buf, iovsDesc, input) {
  const { ptrlens } = iovsDesc
  let written = 0
  for (const [ptr, len] of ptrlens) {
    const extent = Math.min(written + len, input.byteLength)
    buf.set(input.slice(written, extent), ptr)
    written = extent
    if (written === input.byteLength) break
  }
}

/**
 * @this {Runno.WASI}
 * @param {Runno.WASI['fd_read']} origFdRead
 * @returns {Runno.WASI['fd_read']}
 */
function wrapFdRead(origFdRead) {
  return (...args) => {
    const [fd, iovs_ptr, iovs_len, retptr0] = args
    if (fd !== 0) return origFdRead(...args)

    const view = new DataView(this.memory.buffer)
    const iovDescs = collectIOVectors(view, iovs_ptr, iovs_len)

    // not knowing a good reason why the original impl. requests
    // one read per iov

    const input = /** @type {Uint8Array | null} */(
      /** @type {unknown} */(this.context.stdin(iovDescs.lenTotal)))

    if (input == null) {
      return Result.EAGAIN
    }

    const bytes = Math.min(iovDescs.lenTotal, input.byteLength)
    writeIntoIOVectors(new Uint8Array(this.memory.buffer), iovDescs, input)

    // FIXME: missing pushDebugData
    view.setUint32(retptr0, bytes, true)
    return Result.SUCCESS
  }
}

/**
 * @this {Runno.WASI}
 * @param {Runno.WASI['fd_write']} origFdWrite
 * @returns {Runno.WASI['fd_write']}
 */
function wrapFdWrite(origFdWrite) {
  return (...args) => {
    const [fd, ciovs_ptr, ciovs_len, retptr0] = args
    if (fd !== 1 && fd !== 2) return origFdWrite(...args)

    const view = new DataView(this.memory.buffer)
    const iovDescs = collectIOVectors(view, ciovs_ptr, ciovs_len)
    const iov = readIOVectorsMerged(view, iovDescs)

    if (iov.byteLength === 0) {
      return Result.SUCCESS
    }

    const stdfn = fd === 1 ? this.context.stdout : this.context.stderr
    stdfn(/** @type {any} */(iov))

    // FIXME: missing pushDebugData
    view.setUint32(retptr0, iov.byteLength, true)
    return Result.SUCCESS
  }
}

/** @param {Runno.WASI} wasi
 * @param {(timeout: number) => boolean} maybeYieldFunc */
export function patchImportObject(wasi, maybeYieldFunc) {
  const { wasi_snapshot_preview1, ...impObjRest } = wasi.getImportObject()

  const origFdRead = wasi_snapshot_preview1.fd_read
  const origFdWrite = wasi_snapshot_preview1.fd_write
  const origPollOneoff = wasi_snapshot_preview1.poll_oneoff
  return {
    ...impObjRest,
    wasi_snapshot_preview1: {
      ...wasi_snapshot_preview1,
      fd_read: wrapFdRead.bind(wasi)(origFdRead),
      fd_write: wrapFdWrite.bind(wasi)(origFdWrite),
      poll_oneoff: wrapPollOneoff.bind(wasi)(origPollOneoff, maybeYieldFunc)
    }
  }
}

sample.js

import * as Runno from '@runno/wasi'
import { patchImportObject } from './patch-wasi'

/** @typedef {{
 *    stdin: (len: number) => Uint8Array | null,
 *    stdout: (out: Uint8Array) => void,
 *    stderr: (err: Uint8Array) => void,
 *  }} MyStdioDef */

/** @typedef {Partial<
 *    Omit<Runno.WASIContextOptions, 'stdin' | 'stdout' | 'stderr'> & MyStdioDef
 *  >} PatchedRunnoWASIContextOptions */

/** @param {PatchedRunnoWASIContextOptions} opt
 * @returns {Runno.WASIContextOptions} */
function definePatchedRunnoWASIContextOptions(opt) {
  return /** @type {Runno.WASIContextOptions} */(
    /** @type {unknown} */ (opt))
}


const wasi = new WASI(definePatchedRunnoWASIContextOptions({
  stdout(buf) {
    console.log(buf.byteLength)
  }
  /* ... */
})

const importObject = patchImportObject(wasi, timeout => {
  // let waitDuration = timeout < 0 ? Infinity : Math.max(timeout - Date.now(), 0)
  // return stdinReader.pollRead(waitDuration)
})

const wasm = await WebAssembly.instantiateStreaming(fetch('...'), importObject)

wrap-poll-oneoff.js

import * as Runno from '@runno/wasi'
const { Result } = Runno.WASISnapshotPreview1

const EventType = /** @type {const} */({
  CLOCK: 0,
  FD_READ: 1,
  FD_WRITE: 2,
})

const SubscriptionClockFlags = {
  SUBSCRIPTION_CLOCK_ABSTIME: 1,
}

const SUBSCRIPTION_SIZE = 48
const EVENT_SIZE = 32

/** @typedef {{
 *    type: typeof EventType.CLOCK,
 *    id: number, timeout: number, userdata: Uint8Array, precision: number,
 *  }} ClockSubscription */
/** @typedef {{
 *    type: typeof EventType.FD_READ | typeof EventType.FD_WRITE,
 *    fd: number, userdata: Uint8Array,
 *  }} ReadWriteSubscription */

/** @param {Uint8Array} userdata
 * @param {number} error
 * @returns {Uint8Array} */
function createClockEvent(userdata, error) {
  const eventBuffer = new Uint8Array(EVENT_SIZE);
  eventBuffer.set(userdata, 0);

  const view = new DataView(eventBuffer.buffer);
  view.setUint16(8, error, true);
  view.setUint16(10, EventType.CLOCK, true);

  return eventBuffer;
}

/**
 * @this {Runno.WASI}
 * @param {Runno.WASI['poll_oneoff']} origPollOneoff
 * @param {(timeout: number) => boolean} pollStdin
 *   called with -1 or a timeout, should (-1) block or (timeout) return whether stdin is ready
 * @returns {Runno.WASI['poll_oneoff']}
 */
export default function wrapPollOneoff(origPollOneoff, pollStdin) {
  return (...args) => {
    const [in_ptr, out_ptr, nsubscriptions, retptr0] = args

    const subs = []
    for (let i = 0; i < nsubscriptions; i++) {
      const subscriptionBuffer = new Uint8Array(
        this.memory.buffer,
        in_ptr + i * SUBSCRIPTION_SIZE,
        SUBSCRIPTION_SIZE
      );
      subs.push(readSubscription(subscriptionBuffer));
    }

    let stdinIsReady = true

    const readStdinSub = /** @type {ReadWriteSubscription | undefined} */(
      subs.find(s => s.type === EventType.FD_READ && s.fd === 0))
    /** @type {ClockSubscription | undefined} */
    const clockSub = subs.find(s => s.type === EventType.CLOCK)

    // XXX: only handles the two cases that occurs from GHC RTS
    if (readStdinSub) {
      if (subs.length === 1 && clockSub === undefined) {
        // pure (blocking) fd_read
        pollStdin(-1)
      } else if (subs.length === 2 && clockSub !== undefined) {
        // fd_read + clock
        stdinIsReady = pollStdin(clockSub.timeout)
      }
    }  // TODO: handle the case that other fds are queried

    if (!stdinIsReady) {
      // only reports the clock

      const eventBuffer = new Uint8Array(
        this.memory.buffer,
        out_ptr,
        EVENT_SIZE
      );

      eventBuffer.set(
        createClockEvent(/** @type {ClockSubscription} */(clockSub).userdata, Result.SUCCESS)
      )

      const returnView = new DataView(this.memory.buffer, retptr0, 4);
      returnView.setUint32(0, 1, true);
      return Result.SUCCESS
    }

    return origPollOneoff(...args)
  }
}

/** @param {Date} date */
function dateToNanoseconds(date) {
  return BigInt(date.getTime()) * BigInt(1e6);
}

/**
 * @param {Uint8Array} buffer
 * @returns { ReadWriteSubscription | ClockSubscription } */
function readSubscription(buffer) {
  const userdata = new Uint8Array(8);
  userdata.set(buffer.subarray(0, 8));

  const type = buffer[8];

  // View at SubscriptionU offset
  const view = new DataView(buffer.buffer, buffer.byteOffset + 9);
  switch (type) {
    case EventType.FD_READ:
    case EventType.FD_WRITE:
      return {
        userdata,
        type,
        fd: view.getUint32(0, true),
      };
    case EventType.CLOCK:
      const flags = view.getUint16(24, true);
      const currentTimeNanos = dateToNanoseconds(new Date());
      const timeoutRawNanos = view.getBigUint64(8, true);
      const precisionNanos = view.getBigUint64(16, true);

      const timeoutNanos =
        flags & SubscriptionClockFlags.SUBSCRIPTION_CLOCK_ABSTIME
          ? timeoutRawNanos
          : currentTimeNanos + timeoutRawNanos;

      return {
        userdata,
        type,
        id: view.getUint32(0, true),
        timeout: Number(timeoutRawNanos) / 1e6,
        precision: Number(timeoutNanos + precisionNanos) / 1e6,
      };
    default: throw new Error('invalid event type' + type)
  }
}