UDP support

README.md

@@ -7,6 +7,7 @@
 - **Portable:** User-space network stack implemented on top of [`lwIP` + WASM](#why-lwip)
 - **Tun/Tap:** L3 and L2 hooks using virtual [`TunInterface`](#tun-interface) and [`TapInterface`](#tap-interface)
 - **TCP API:** Establish TCP connections over the virtual network stack using [clients](#connecttcp) and [servers](#listentcp)
+- **UDP API:** Send and receive UDP datagrams over the virtual network stack using [sockets](#openudp)
 - **Cross platform**: Built on web standard APIs (`ReadableStream`, `WritableStream`, etc)
 - **Lightweight:** Less than 100KB
 - **Fast:** Over 500Mbps between stacks
@@ -465,6 +466,102 @@ To close the connection, call `close()`:
 await connection.close();
 ```
 
+## UDP API
+
+The UDP API allows you to send and receive UDP datagrams over the virtual network stack.
+
+### `openUdp()`
+
+To open a UDP socket, call `openUdp()`:
+
+```ts
+const udpSocket = await stack.openUdp();
+```
+
+Since UDP is connectionless, `openUdp()` is used to create a socket that can both listen for UDP datagrams and send UDP datagrams. It returns a [`UdpSocket`](#udpsocket) that you can use to send and receive data.
+
+Passing no arguments to `openUdp()` will create a socket that sends and receives datagrams on any interface (ie. `0.0.0.0`) and on a random port. If you want to bind to a specific IP address or port, you can pass an options object:
+
+```ts
+const udpSocket = await stack.openUdp({
+  ip: '10.0.0.1',
+  port: 1234,
+});
+```
+
+If you are creating a UDP server, you would typically just bind to a port:
+
+```ts
+const udpSocket = await stack.openUdp({
+  port: 1234,
+});
+```
+
+If you are creating a UDP client, you would typically let the stack choose a random port:
+
+```ts
+const udpSocket = await stack.openUdp();
+```
+
+### `UdpSocket`
+
+A `UdpSocket` represents a bound UDP socket. It exposes a [`ReadableStream`](https://developer.mozilla.org/docs/Web/API/ReadableStream) and [`WritableStream`](https://developer.mozilla.org/docs/Web/API/WritableStream) as the underlying APIs to send and receive data. It also implements the async iterable protocol for convenience.
+
+```ts
+type UdpDatagram = {
+  host: string;
+  port: number;
+  data: Uint8Array;
+};
+
+interface UdpSocket {
+  readable: ReadableStream<UdpDatagram>;
+  writable: WritableStream<UdpDatagram>;
+  close(): Promise<void>;
+  [Symbol.asyncIterator](): AsyncIterator<UdpDatagram>;
+}
+```
+
+You would typically read incoming datagrams by iterating over the `UdpSocket` using the `for await` syntax:
+
+```ts
+for await (const datagram of udpSocket) {
+  console.log(
+    datagram.host,
+    datagram.port,
+    new TextDecoder().decode(datagram.data)
+  );
+}
+```
+
+Notice that each datagram is an object with `host`, `port`, and `data` properties. This is because UDP is connectionless so we need a way to identify the sender of each datagram.
+
+You can also read datagrams from the `readable` stream directly by acquiring a reader:
+
+```ts
+const reader = udpSocket.readable.getReader();
+
+// Read datagram
+const { value, done } = await reader.read();
+```
+
+To send datagrams, you would typically acquire a writer from the `writable` stream:
+
+```ts
+const writer = udpSocket.writable.getWriter();
+
+// Send datagram
+await writer.write({
+  host: '10.0.0.2',
+  port: 1234,
+  data: new TextEncoder().encode('Hello, world!'),
+});
+```
+
+Outbound datagrams follow the same format as inbound datagrams: an object with `host`, `port`, and `data` properties indicating the destination host, port, and data.
+
+Unlike Tun and Tap interfaces which are also connectionless, UDP sockets do not require you to lock the stream before receiving data - simply calling `stack.openUdp()` will begin listening for datagrams.
+
 ## Frameworks/bundlers
 
 Some frameworks require additional configuration to correctly load WASM files (which `tcpip.js` depends on). Here are some common frameworks and how to configure them:
@@ -489,7 +586,6 @@ _Background:_ Vite optimizes dependencies during development to improve build ti
 ## Future plans
 
 - [ ] HTTP API
-- [ ] UDP API
 - [ ] ICMP (ping) API
 - [ ] DHCP API
 - [ ] DNS API

packages/tcpip/src/bindings/udp.ts

@@ -0,0 +1,186 @@
+import { LwipError } from '../lwip/errors.js';
+import {
+  parseIPv4Address,
+  serializeIPv4Address,
+  type IPv4Address,
+} from '../protocols/ipv4.js';
+import type { Pointer } from '../types.js';
+import { EventMap, fromReadable, Hooks, nextMicrotask } from '../util.js';
+import { Bindings } from './base.js';
+
+export type UdpDatagram = {
+  host: IPv4Address;
+  port: number;
+  data: Uint8Array;
+};
+
+type UdpSocketHandle = Pointer;
+
+type UdpSocketOuterHooks = {
+  send(datagram: UdpDatagram): Promise<void>;
+  close(): Promise<void>;
+};
+
+type UdpSocketInnerHooks = {
+  receive(datagram: UdpDatagram): Promise<void>;
+};
+
+const UdpSocketHooks = new Hooks<
+  UdpSocket,
+  UdpSocketOuterHooks,
+  UdpSocketInnerHooks
+>();
+
+export type UdpImports = {
+  receive_udp_datagram(
+    handle: UdpSocketHandle,
+    ip: number,
+    port: number,
+    datagramPtr: number,
+    length: number
+  ): Promise<void>;
+};
+
+export type UdpExports = {
+  open_udp_socket(host: Pointer | null, port: number): UdpSocketHandle;
+  close_udp_socket(handle: UdpSocketHandle): void;
+  send_udp_datagram(
+    handle: UdpSocketHandle,
+    ip: Pointer | null,
+    port: number,
+    datagram: Pointer,
+    length: number
+  ): number;
+};
+
+export class UdpBindings extends Bindings<UdpImports, UdpExports> {
+  #UdpSockets = new EventMap<UdpSocketHandle, UdpSocket>();
+
+  imports = {
+    receive_udp_datagram: async (
+      handle: UdpSocketHandle,
+      hostPtr: number,
+      port: number,
+      datagramPtr: number,
+      length: number
+    ) => {
+      const host = this.copyFromMemory(hostPtr, 4);
+      const datagram = this.copyFromMemory(datagramPtr, length);
+      const socket = this.#UdpSockets.get(handle);
+
+      if (!socket) {
+        console.error('received datagram on unknown udp socket');
+        return;
+      }
+
+      // Wait for synchronous lwIP operations to complete to prevent reentrancy issues
+      await nextMicrotask();
+
+      UdpSocketHooks.getInner(socket).receive({
+        host: parseIPv4Address(host),
+        port,
+        data: datagram,
+      });
+    },
+  };
+
+  async open(options: UdpSocketOptions) {
+    using hostPtr = options.host
+      ? this.copyToMemory(serializeIPv4Address(options.host))
+      : null;
+
+    const handle = this.exports.open_udp_socket(hostPtr, options.port ?? 0);
+
+    const udpSocket = new UdpSocket();
+
+    UdpSocketHooks.setOuter(udpSocket, {
+      send: async (datagram: UdpDatagram) => {
+        using hostPtr = this.copyToMemory(serializeIPv4Address(datagram.host));
+        using datagramPtr = this.copyToMemory(datagram.data);
+
+        const result = this.exports.send_udp_datagram(
+          handle,
+          hostPtr,
+          datagram.port,
+          datagramPtr,
+          datagram.data.length
+        );
+
+        if (result !== LwipError.ERR_OK) {
+          throw new Error(`failed to send udp datagram: ${result}`);
+        }
+      },
+      close: async () => {
+        this.exports.close_udp_socket(handle);
+        this.#UdpSockets.delete(handle);
+      },
+    });
+
+    this.#UdpSockets.set(handle, udpSocket);
+
+    return udpSocket;
+  }
+}
+
+export type UdpSocketOptions = {
+  /**
+   * The local host to bind to.
+   *
+   * If not provided, the socket will bind to all available interfaces.
+   */
+  host?: IPv4Address;
+
+  /**
+   * The local port to bind to.
+   *
+   * If not provided, the socket will bind to a random port.
+   */
+  port?: number;
+};
+
+export class UdpSocket implements AsyncIterable<UdpDatagram> {
+  #readableController?: ReadableStreamDefaultController<UdpDatagram>;
+  #writableController?: WritableStreamDefaultController;
+
+  readable: ReadableStream<UdpDatagram>;
+  writable: WritableStream<UdpDatagram>;
+
+  constructor() {
+    UdpSocketHooks.setInner(this, {
+      receive: async (datagram: UdpDatagram) => {
+        if (!this.#readableController) {
+          throw new Error('readable controller not initialized');
+        }
+        this.#readableController.enqueue(datagram);
+      },
+    });
+
+    this.readable = new ReadableStream({
+      start: (controller) => {
+        this.#readableController = controller;
+      },
+    });
+
+    this.writable = new WritableStream({
+      start: (controller) => {
+        this.#writableController = controller;
+      },
+      write: async (datagram) => {
+        await UdpSocketHooks.getOuter(this).send(datagram);
+      },
+    });
+  }
+
+  async close() {
+    await UdpSocketHooks.getOuter(this).close();
+    this.#readableController?.error(new Error('udp socket closed'));
+    this.#writableController?.error(new Error('udp socket closed'));
+  }
+
+  [Symbol.asyncIterator](): AsyncIterator<UdpDatagram> {
+    if (this.readable.locked) {
+      throw new Error('readable stream already locked');
+    }
+    return fromReadable(this.readable);
+  }
+}

packages/tcpip/src/index.ts

@@ -1,4 +1,4 @@
-export { NetworkStack, createStack } from './stack.js';
+export { createStack, NetworkStack } from './stack.js';
 export type { NetworkInterface } from './types.js';
 
 export {
@@ -22,3 +22,9 @@ export {
   type TcpConnectionOptions,
   type TcpListenerOptions,
 } from './bindings/tcp.js';
+
+export {
+  UdpSocket,
+  type UdpDatagram,
+  type UdpSocketOptions,
+} from './bindings/udp.js';