fix(net): never leave a torn request message on the stream

js/net/src/ietf/message.ts

@@ -1,4 +1,4 @@
-import { concat, encodeVarint, Reader, Writer } from "../stream.ts";
+import { concat, encodeVarint, Reader, type StreamVersion, Writer } from "../stream.ts";
 import { MalformedAuthToken } from "./token.ts";
 import type { IetfVersion } from "./version.ts";
 
@@ -23,8 +23,13 @@ export interface Encode {
 // Writes `msg` after the type `id` the caller picked for it, in one write, so a message that
 // fails to encode leaves nothing on the stream.
 export async function write(writer: Writer, id: number, msg: Encode, version: IetfVersion) {
-	const body = await Writer.collect(writer.version, (w) => msg.encode(w, version));
-	await writer.write(concat([encodeVarint(id, writer.version), body]));
+	await writer.write(await frame(writer.version, id, msg, version));
+}
+
+// The one chunk `write` sends, for a caller that must tell a failed encode from a failed write.
+export async function frame(stream: StreamVersion, id: number, msg: Encode, version: IetfVersion) {
+	const body = await Writer.collect(stream, (w) => msg.encode(w, version));
+	return concat([encodeVarint(id, stream), body]);
 }
 
 // Reads a message with a u16 size prefix.

js/net/src/ietf/subscriber.ts

@@ -13,6 +13,7 @@ import {
 	error,
 	ProtocolViolation,
 	reason,
+	SessionCode,
 	StreamCode,
 	Stream as StreamError,
 	sessionCause,
@@ -466,11 +467,11 @@ export class Subscriber {
 		}
 	}
 
-	// The peer broke the protocol: close with the violation's code before the session's own
-	// close, which would send a clean one, as the connection does for its own violations.
-	#violated(err: ProtocolViolation) {
+	// Close with this code before the session's own close, which would send a clean one, as
+	// the connection does for its own violations.
+	#fail(code: SessionCode, reason: string) {
 		try {
-			this.#quic?.close({ closeCode: err.code, reason: err.message });
+			this.#quic?.close({ closeCode: code, reason });
 		} catch {
 			// Already closed.
 		}
@@ -1062,10 +1063,12 @@ export class Subscriber {
 		void stopped.then(() => stop.abort());
 
 		for (;;) {
+			// With no subscriber left there is nothing to follow: the subscription is ending, or
+			// keeps what it last asked for until one returns, as on lite and in Rust.
 			const subscription = producer.subscription.peek();
 			const wanted = toWire(subscription?.priority ?? 0);
-			const wantedMaxAge = this.#maxAge ? (subscription?.maxDelay ?? Number.MAX_SAFE_INTEGER) : undefined;
-			if ((wanted === sent && wantedMaxAge === sentMaxAge) || unanswered.peek() > 0) {
+			const wantedMaxAge = this.#maxAge ? subscription?.maxDelay : undefined;
+			if (!subscription || (wanted === sent && wantedMaxAge === sentMaxAge) || unanswered.peek() > 0) {
 				const disposes: Dispose[] = [];
 				const wake = new Promise<void>((resolve) => {
 					disposes.push(
@@ -1093,10 +1096,24 @@ export class Subscriber {
 			const update = legacy
 				? new SubscribeUpdate({ requestId, ownRequestId: own, parameters })
 				: new SubscribeUpdate({ requestId: own, parameters });
+			let message: Uint8Array;
+			try {
+				message = await Message.frame(stream.writer.version, SubscribeUpdate.id, update, version);
+			} catch (err) {
+				// Nothing reached the wire, but the track can no longer follow its subscribers.
+				console.warn(`subscribe update failed: id=${requestId} error=${reason(error(err))}`);
+				producer.close(error(err));
+				// Drafts 14 and 15 require the next request ID, so one spent and never sent would
+				// stall every later request: as in Rust, that ends the session.
+				if (version === Version.DRAFT_14 || version === Version.DRAFT_15) {
+					this.#fail(SessionCode.Internal, "subscribe update failed to encode");
+				}
+				return;
+			}
 			// Counted before writing: the answer can be read before the write's promise resolves.
 			if (answered) unanswered.update((count) => count + 1);
 			try {
-				await Message.write(stream.writer, SubscribeUpdate.id, update, version);
+				await stream.writer.write(message);
 			} catch {
 				// The publisher stopped reading; its responses decide how the track ends.
 				return;
@@ -1504,7 +1521,7 @@ export class Subscriber {
 			// The control message establishing this alias can arrive after the data stream.
 			subscription = await this.#aliases.get(group.trackAlias);
 		} catch (err: unknown) {
-			if (err instanceof ProtocolViolation) this.#violated(err);
+			if (err instanceof ProtocolViolation) this.#fail(err.code, err.message);
 			const e = await sessionCause(this.#quic, err);
 			// Ours: we cancelled the subscription and the publisher has not stopped yet.
 			// Anything else on this alias is the publisher sending data for a track it never
@@ -1602,7 +1619,7 @@ export class Subscriber {
 			// A group with no objects still exists.
 			open().close();
 		} catch (err: unknown) {
-			if (err instanceof ProtocolViolation) this.#violated(err);
+			if (err instanceof ProtocolViolation) this.#fail(err.code, err.message);
 			const e = await sessionCause(this.#quic, err);
 			// The compact form names the subgroup by its first Object ID, so a nonzero one is
 			// subgroup N: stop the stream before any group exists for it.

js/net/src/subscribe_priority.test.ts

@@ -2,15 +2,20 @@
 // follows it as they come and go. The JS origin never serves a received route onward, so the
 // relay case lives with the Rust relay (rs/moq-net/tests/integration/subscribe_priority.rs).
 
-import { expect, test } from "bun:test";
+import { expect, spyOn, test } from "bun:test";
+import { Manual } from "@moq/clock";
 import { accept, connect, type Established } from "./connection/index.ts";
+import { SessionCode, type Session as SessionError } from "./error.ts";
 import * as Ietf from "./ietf/index.ts";
+import { SubscribeUpdate } from "./ietf/subscribe.ts";
 import * as Lite from "./lite/index.ts";
 import { createMockTransportPair, type MockTransport } from "./mock.ts";
 import { Producer as OriginProducer } from "./origin.ts";
 import * as Path from "./path.ts";
+import { Writer } from "./stream.ts";
 import { Timescale } from "./time.ts";
 import type { Producer as TrackProducer } from "./track.ts";
+import { nextTask } from "./util/yield.ts";
 import { wireOf } from "./wire.ts";
 
 const url = new URL("https://localhost:4443/test");
@@ -89,3 +94,115 @@ test.each([Lite.ALPN_07_WIP, ...Object.values(Ietf.ALPN)])(
 		server.abort();
 	},
 );
+
+// A write may yield to the host after it lands (the shared write budget does under load), so the
+// peer's end of a request can arrive between any two writes. With every write yielding, the last
+// subscribers leaving while an update is answered must leave only whole messages on the stream:
+// a torn one closes the session as malformed.
+test.each(Object.values(Ietf.ALPN))("%s ends a subscription mid-update without a torn message", async (protocol) => {
+	const clock = new Manual();
+	const installed = clock.install();
+	const write = Writer.prototype.write;
+	const yielding = spyOn(Writer.prototype, "write").mockImplementation(async function (this: Writer, chunk) {
+		await write.call(this, chunk);
+		await nextTask();
+	});
+	const origin = new OriginProducer();
+	const broadcast = origin.createBroadcast(path);
+	broadcast.announce();
+	const video = broadcast.createTrack("video", { timescale: Timescale.MILLI });
+	const [client, server] = await session(protocol, origin);
+	let closed: Error | null | undefined;
+	void client.closed.then((err) => {
+		closed = err;
+	});
+	try {
+		const remote = wireOf(client).consume(path);
+		const low = remote.track("video").subscribe({ priority: LOW });
+		await priority(video, LOW);
+		const high = remote.track("video").subscribe({ priority: HIGH });
+		await priority(video, HIGH);
+		high.close();
+		await priority(video, LOW);
+		// The publisher has read the update and is answering it as the subscription ends.
+		low.close();
+		while (video.subscription.peek() !== undefined) await video.subscription.changed();
+		// Each hop is a task or two away, so this lets the subscriber read all the publisher wrote.
+		for (let i = 0; i < 20; i++) await nextTask();
+		expect(closed).toBeUndefined();
+		remote.close();
+	} finally {
+		yielding.mockRestore();
+		broadcast.close();
+		client.abort();
+		server.abort();
+		installed.close();
+	}
+});
+
+// With no subscriber left there is no aggregate to follow, so the subscription ends without an
+// update of defaults first, as on lite and in Rust.
+test.each(Object.values(Ietf.ALPN))("%s sends no update when the last subscriber leaves", async (protocol) => {
+	const origin = new OriginProducer();
+	const broadcast = origin.createBroadcast(path);
+	broadcast.announce();
+	const video = broadcast.createTrack("video", { timescale: Timescale.MILLI });
+	const [client, server] = await session(protocol, origin);
+	try {
+		const remote = wireOf(client).consume(path);
+		const track = remote.track("video").subscribe({ priority: LOW });
+		await priority(video, LOW);
+		await track.info();
+		track.close();
+		while (video.subscription.peek() !== undefined) {
+			expect(video.subscription.peek()?.priority).toBe(LOW);
+			await video.subscription.changed();
+		}
+		remote.close();
+	} finally {
+		broadcast.close();
+		client.abort();
+		server.abort();
+	}
+});
+
+// An update that fails to encode is a bug, not a publisher that stopped reading, so it ends the
+// track with the error rather than leaving it at a stale priority. Nothing reaches the wire.
+// Drafts 14 and 15 require the next request ID, so the one the update spent ends the session.
+test.each(Object.values(Ietf.ALPN))("%s ends the track when an update fails to encode", async (protocol) => {
+	const origin = new OriginProducer();
+	const broadcast = origin.createBroadcast(path);
+	broadcast.announce();
+	const video = broadcast.createTrack("video", { timescale: Timescale.MILLI });
+	const [client, server] = await session(protocol, origin);
+	let closed: Error | null | undefined;
+	void client.closed.then((err) => {
+		closed = err;
+	});
+	const broken = new Error("broken update");
+	const encode = spyOn(SubscribeUpdate.prototype, "encode").mockRejectedValue(broken);
+	try {
+		const remote = wireOf(client).consume(path);
+		const track = remote.track("video").subscribe({ priority: LOW });
+		await priority(video, LOW);
+		await track.info();
+		const high = remote.track("video").subscribe({ priority: HIGH });
+		expect(await track.closed).toBe(broken);
+		expect(await high.closed).toBe(broken);
+		while (video.subscription.peek() !== undefined) {
+			expect(video.subscription.peek()?.priority).toBe(LOW);
+			await video.subscription.changed();
+		}
+		if (protocol === Ietf.ALPN.DRAFT_14 || protocol === Ietf.ALPN.DRAFT_15) {
+			expect(((await server.closed) as SessionError | null)?.code).toBe(SessionCode.Internal);
+		} else {
+			expect(closed).toBeUndefined();
+		}
+		remote.close();
+	} finally {
+		encode.mockRestore();
+		broadcast.close();
+		client.abort();
+		server.abort();
+	}
+});

quest/m1/test-flakes-2/README.md

@@ -39,7 +39,6 @@ Public API: none. Wire: none.
 
 ## Required
 
-- [Union renewal](/quest/m1/test-flakes-2/auth-union-renewal.md) - the JS subscriber never leaves a torn REQUEST_UPDATE on the stream
 - [Media resource instrumentation](/quest/m1/test-flakes-2/resource-instrumentation.md) - the media test stops seeing a playing player report nothing to release
 - [Tone check grace](/quest/m1/test-flakes-2/tone-grace.md) - the interop tone check charges its grace only for audible playback, so a refilling ring under CPU starvation is not a failure
 - [Relay tests in one process](/quest/m1/test-flakes-2/relay-one-process.md) - moq-relay's integration binary passes under plain `cargo test`, not only nextest

rs/moq-net/src/coding/writer.rs

@@ -243,10 +243,27 @@ impl<S: crate::transport::poll::SendStream, V: StreamCodes> Writer<S, V> {
 }
 
 impl<S: crate::transport::poll::SendStream> Writer<S, ietf::Version> {
-	/// Encode an IETF `Message` to the stream, writing `[type_id][size][body]`.
+	/// Encode `msg` into the write buffer as `[id][size][body]`, to be sent by
+	/// [`Self::poll_flush`]. The type and body go in together or not at all, so a failed encode
+	/// never leaves a bare type for the next write to complete. `id` is explicit for the drafts
+	/// that send one body under several types.
+	pub fn buffer_message<T: Encode<ietf::Version>>(&mut self, id: u64, msg: &T) -> Result<(), Error> {
+		let start = self.buffer.len();
+		self.buffer_varint(id)?;
+		// The body is buffered on its own so control stats count it without its type selector.
+		let version = self.version;
+		if let Err(err) = self.buffer_with(|w| msg.encode(w, version)) {
+			self.buffer.truncate(start);
+			return Err(err);
+		}
+		Ok(())
+	}
+
+	/// Encode an IETF `Message` to the stream as `[type_id][size][body]`, with the same
+	/// all-or-nothing buffering as [`Self::buffer_message`] and cancellation rules as [`Self::encode`].
 	pub async fn encode_message<T: ietf::Message>(&mut self, msg: &T) -> Result<(), Error> {
-		self.buffer_varint(T::ID)?;
-		self.encode(msg).await
+		self.buffer_message(T::ID, msg)?;
+		std::future::poll_fn(|cx| self.poll_flush(cx)).await
 	}
 }
 
@@ -357,8 +374,8 @@ mod tests {
 	#[derive(Debug)]
 	struct Poison;
 
-	impl Encode<crate::lite::Version> for Poison {
-		fn encode(&self, w: &mut Encoder<'_>, _: crate::lite::Version) -> Result<(), EncodeError> {
+	impl<V> Encode<V> for Poison {
+		fn encode(&self, w: &mut Encoder<'_>, _: V) -> Result<(), EncodeError> {
 			w.control = Some(crate::stats::Control::Probe);
 			w.slice(b"junk");
 			Err(EncodeError::BoundsExceeded)
@@ -379,6 +396,17 @@ mod tests {
 		let mut cx = std::task::Context::from_waker(Waker::noop());
 		assert!(writer.poll_flush(&mut cx).is_ready());
 		assert_eq!(*log.writes.lock().unwrap(), vec![5, 7], "the partial encode leaked");
+
+		// A control message's type goes with its body: a bare type would make the next
+		// message, or a FIN, a torn one.
+		let mut writer = Writer::new(SinkSend::new(Log::default()), crate::ietf::Version::Draft17);
+		writer.buffer_varint(5).unwrap();
+		writer.buffer_message(0x2, &Poison).unwrap_err();
+		writer.buffer_varint(7).unwrap();
+
+		let log = writer.stream.as_ref().unwrap().log.clone();
+		assert!(writer.poll_flush(&mut cx).is_ready());
+		assert_eq!(*log.writes.lock().unwrap(), vec![5, 7], "the message type leaked");
 	}
 
 	#[test]

rs/moq-net/src/ietf/adapter.rs

@@ -2217,8 +2217,7 @@ mod tests {
 				request_id: RequestId(0),
 			};
 			let size = msg.encode_bytes(version).unwrap().len() as u64;
-			virtual_stream.writer.buffer_varint(ietf::FetchCancel::ID).unwrap();
-			virtual_stream.writer.encode(&msg).await.unwrap();
+			virtual_stream.writer.encode_message(&msg).await.unwrap();
 			let raw = adapter.shared.control.pop().await.unwrap();
 			let log = Log::default();
 			let mut physical = Writer::new(