1// A websocket client for the Durable Object tests: NIP-01 over an upgrade2// through the Worker, with the frames queued so a test reads them in order.3import { SELF } from "cloudflare:test";4import { expect } from "vitest";5import type { Event } from "nostr-tools/pure";6import { ev } from "./relay.ts";78export class WS {9 private queue: any[][] = [];10 private waiters: ((m: any[]) => void)[] = [];11 private n = 0;12 challenge = "";13 closed: { code: number; reason: string } | null = null;14 frames: string[] = []; // every frame received, for leak checks15 constructor(public ws: WebSocket) {16 ws.accept();17 ws.addEventListener("message", (e) => {18 this.frames.push(e.data as string);19 const m = JSON.parse(e.data as string);20 const w = this.waiters.shift();21 if (w) w(m);22 else this.queue.push(m);23 });24 ws.addEventListener("close", (e) => {25 this.closed = { code: e.code, reason: e.reason };26 });27 }28 // connect opens a socket, from an address when one is given, and reads the29 // AUTH challenge.30 static async connect(host: string, ip?: string): Promise<WS> {31 const c = await WS.tryConnect(host, ip);32 if (!c) throw new Error(`${host} refused the upgrade`);33 return c;34 }35 // tryConnect is connect for a socket the relay may refuse: null then.36 static async tryConnect(host: string, ip?: string): Promise<WS | null> {37 const headers: Record<string, string> = { upgrade: "websocket" };38 if (ip) headers["cf-connecting-ip"] = ip;39 const resp = await SELF.fetch(`http://${host}/`, { headers });40 if (!resp.webSocket) return null;41 const c = new WS(resp.webSocket);42 c.challenge = (await c.expect("AUTH"))[1];43 return c;44 }45 send(...m: unknown[]) {46 this.ws.send(JSON.stringify(m));47 }48 raw(s: string) {49 this.ws.send(s);50 }51 recv(): Promise<any[]> {52 const m = this.queue.shift();53 if (m) return Promise.resolve(m);54 return new Promise((res) => this.waiters.push(res));55 }56 // recvOr resolves null when nothing arrives within ms.57 recvOr(ms: number): Promise<any[] | null> {58 const m = this.queue.shift();59 if (m) return Promise.resolve(m);60 return new Promise((res) => {61 this.waiters.push(res);62 setTimeout(() => {63 const i = this.waiters.indexOf(res);64 if (i >= 0) {65 this.waiters.splice(i, 1);66 res(null);67 }68 }, ms);69 });70 }71 async expect(type: string) {72 const m = await this.recv();73 expect(m[0], JSON.stringify(m)).toBe(type);74 return m;75 }76 // ok publishes and waits for the OK. Live pushes the relay fans out while77 // handling the event (its own put-user, say) arrive first; they are kept78 // in order for a later expect.79 async ok(e: Event) {80 this.send("EVENT", e);81 const held: any[][] = [];82 for (;;) {83 const m = await this.recv();84 if (m[0] === "OK" && m[1] === e.id) {85 this.queue.unshift(...held);86 return { ok: m[2] as boolean, msg: m[3] as string };87 }88 held.push(m);89 }90 }91 // auth proves a key over the socket.92 async auth(sk: Uint8Array, host: string) {93 this.send("AUTH", ev(sk, 22242, "", [["relay", "ws://" + host], ["challenge", this.challenge]]));94 const m = await this.expect("OK");95 expect(m[2], m[3]).toBe(true);96 }97 // query is a one-shot subscription: unique id, closed after EOSE so it98 // never leaks live pushes into a later query. The events and "" on EOSE,99 // or the CLOSED reason.100 async query(filter: unknown) {101 const id = "q" + ++this.n;102 const r = await this.open(id, filter);103 if (!r.closed) this.send("CLOSE", id);104 return r;105 }106 // req is query for a subscription that must be accepted: the events.107 async req(filter: unknown): Promise<Event[]> {108 const r = await this.query(filter);109 if (r.closed) throw new Error(r.closed);110 return r.events;111 }112 // open subscribes under a given id and stays subscribed: the stored events113 // and "" on EOSE, or the CLOSED reason.114 async open(id: string, ...filters: unknown[]) {115 this.send("REQ", id, ...filters);116 const events: Event[] = [];117 const held: any[][] = []; // frames for other subscriptions, kept in order118 const done = (closed: string) => {119 this.queue.unshift(...held);120 return { events, closed };121 };122 for (;;) {123 const m = await this.recv();124 if (m[0] === "EVENT" && m[1] === id) events.push(m[2]);125 else if (m[0] === "EOSE" && m[1] === id) return done("");126 else if (m[0] === "CLOSED" && m[1] === id) return done(m[2] as string);127 else if (m[0] === "EVENT" || m[0] === "CLOSED") held.push(m);128 else throw new Error(JSON.stringify(m));129 }130 }131 // count asks for a COUNT: the number and "" or the CLOSED reason.132 async count(filter: unknown, id = "c" + ++this.n) {133 this.send("COUNT", id, filter);134 const m = await this.recv();135 return m[0] === "COUNT" ? { count: m[2].count as number, closed: "" } : { count: -1, closed: m[2] as string };136 }137 // sync opens a NIP-77 reconciliation and returns the NEG-ERR reason, or "".138 async sync(id: string, filter: unknown) {139 this.send("NEG-OPEN", id, filter, "61");140 const m = await this.recv();141 return m[0] === "NEG-ERR" ? (m[2] as string) : "";142 }143}