diff --git a/ios/ClawChat/Sources/ClawChat/Conversation.swift b/ios/ClawChat/Sources/ClawChat/Conversation.swift index da0ca12..4bb436a 100644 --- a/ios/ClawChat/Sources/ClawChat/Conversation.swift +++ b/ios/ClawChat/Sources/ClawChat/Conversation.swift @@ -31,6 +31,15 @@ public struct ChatAttachment: Equatable, Sendable { self.filename = filename self.mime = mime } + + /// The wire-level kind this attachment travels as. + public var wireKind: Wire.MediaKind { + switch kind { + case .image: return .image + case .audio: return .audio + case .file: return .file + } + } } public enum ChatDeliveryState: Equatable, Sendable { @@ -237,6 +246,27 @@ public final class Conversation: ObservableObject { messages[idx] = msg saveToStore() } + // A message carrying an attachment has to go back out over the + // media path — the text path would transmit only its caption. + if let attachment = m.attachments.first { + let caption = m.text.isEmpty ? nil : m.text + Task { [weak self] in + do { + _ = try await conn.send( + media: attachment.wireKind, + bytes: attachment.bytes, + filename: attachment.filename, + mime: attachment.mime, + caption: caption, + messageId: m.id + ) + self?.updateDelivery(id: m.id, to: .sent) + } catch { + self?.updateDelivery(id: m.id, to: deliveryStateForSendError(error)) + } + } + continue + } Task { [weak self] in do { _ = try await conn.send(text: m.text, messageId: m.id) @@ -299,24 +329,25 @@ public final class Conversation: ObservableObject { filename: String? = nil, mime: String? = nil ) { - guard let conn = connection, conn.isReady else { return } let caption = draft.trimmingCharacters(in: .whitespacesAndNewlines) draft = "" let messageId = Wire.newId() + let attachment = ChatAttachment(kind: kind, bytes: bytes, filename: filename, mime: mime) let chat = ChatMessage( id: messageId, sender: .me, text: caption, - attachments: [ChatAttachment(kind: kind, bytes: bytes, filename: filename, mime: mime)], + attachments: [attachment], delivery: .sending ) + // Append and persist before touching the connection, mirroring send(). + // The message survives an app restart while still in .sending state. messages.append(chat) - let wireKind: Wire.MediaKind - switch kind { - case .image: wireKind = .image - case .audio: wireKind = .audio - case .file: wireKind = .file - } + saveToStore() + // If we're not connected, leave it as .sending — drainOutbox() will + // pick it up the moment the connection becomes ready. + guard let conn = connection, conn.isReady else { return } + let wireKind = attachment.wireKind Task { [weak self] in do { _ = try await conn.send( diff --git a/ios/ClawChat/Tests/ClawChatTests/ConversationOutboxTests.swift b/ios/ClawChat/Tests/ClawChatTests/ConversationOutboxTests.swift index 3befd58..2b40513 100644 --- a/ios/ClawChat/Tests/ClawChatTests/ConversationOutboxTests.swift +++ b/ios/ClawChat/Tests/ClawChatTests/ConversationOutboxTests.swift @@ -98,4 +98,87 @@ final class ConversationOutboxTests: XCTestCase { XCTAssertEqual(c.messages.count, 1) XCTAssertEqual(c.messages.first?.delivery, .sending) } + + // MARK: - Attachments queued while offline + + func testOfflineAttachmentIsQueuedAsSending() { + // sendAttachment with no connection must behave like send(): append + // the message in .sending rather than discarding it. + let c = Conversation() + c.messageStore = MessageStore(profileDir: workDir) + c.draft = "look at this" + c.sendAttachment( + kind: .image, + bytes: Data([0xFF, 0xD8, 0xFF, 0xE0]), + filename: "cat.jpg", + mime: "image/jpeg" + ) + XCTAssertEqual(c.messages.count, 1) + let m = try? XCTUnwrap(c.messages.first) + XCTAssertEqual(m?.delivery, .sending) + XCTAssertEqual(m?.sender, .me) + XCTAssertEqual(m?.text, "look at this", "caption should ride along with the attachment") + XCTAssertEqual(m?.attachments.count, 1) + XCTAssertEqual(m?.attachments.first?.filename, "cat.jpg") + XCTAssertEqual(c.draft, "", "draft is consumed as the caption") + } + + func testOfflineAttachmentSurvivesAppRestartWithBytesIntact() { + let bytes = Data([0x01, 0x02, 0x03, 0x04, 0x05]) + do { + let c = Conversation() + c.messageStore = MessageStore(profileDir: workDir) + c.sendAttachment(kind: .file, bytes: bytes, filename: "doc.bin", mime: "application/octet-stream") + XCTAssertEqual(c.messages.count, 1) + } + + // Fresh Conversation + store on the same dir, as after a cold launch. + let c2 = Conversation() + c2.messageStore = MessageStore(profileDir: workDir) + c2.loadFromStoreIfAvailable() + XCTAssertEqual(c2.messages.count, 1, "the queued attachment must be persisted, not dropped") + let m = c2.messages.first + XCTAssertEqual(m?.delivery, .sending) + XCTAssertEqual(m?.attachments.count, 1, "attachment must survive the round-trip") + XCTAssertEqual(m?.attachments.first?.bytes, bytes, "attachment bytes must be byte-identical") + XCTAssertEqual(m?.attachments.first?.kind, .file) + XCTAssertEqual(m?.attachments.first?.mime, "application/octet-stream") + } + + func testRehydratedAttachmentIsSelectedByTheDrainQueue() { + // drainOutbox re-sends every from-me message left in .sending/.failed. + // A rehydrated media message must qualify, and must still carry the + // attachment that the media send path needs. + let c = Conversation() + c.messageStore = MessageStore(profileDir: workDir) + c.draft = "caption" + c.sendAttachment(kind: .audio, bytes: Data([0x11, 0x22]), filename: "vm.m4a", mime: "audio/mp4") + c.draft = "text only" + c.send() + + let c2 = Conversation() + c2.messageStore = MessageStore(profileDir: workDir) + c2.loadFromStoreIfAvailable() + + let pending = c2.messages.filter { $0.sender == .me && $0.delivery == .sending } + XCTAssertEqual(pending.count, 2) + let withMedia = pending.filter { !$0.attachments.isEmpty } + XCTAssertEqual(withMedia.count, 1, "the media message must be part of the drain set") + XCTAssertEqual(withMedia.first?.attachments.first?.kind, .audio) + // The drain branches on this to pick the media send path over text. + XCTAssertEqual(withMedia.first?.attachments.first?.wireKind, .audio) + XCTAssertEqual(withMedia.first?.text, "caption") + + // And the text-only message must still route via the text path. + let textOnly = pending.filter { $0.attachments.isEmpty } + XCTAssertEqual(textOnly.count, 1) + XCTAssertEqual(textOnly.first?.text, "text only") + } + + func testAttachmentWireKindMapping() { + // The mapping drainOutbox uses to re-send a persisted attachment. + XCTAssertEqual(ChatAttachment(kind: .image, bytes: Data()).wireKind, .image) + XCTAssertEqual(ChatAttachment(kind: .audio, bytes: Data()).wireKind, .audio) + XCTAssertEqual(ChatAttachment(kind: .file, bytes: Data()).wireKind, .file) + } } diff --git a/ios/ClawChat/Tests/ClawChatTests/ConversationTests.swift b/ios/ClawChat/Tests/ClawChatTests/ConversationTests.swift index 385dd6e..d5fa737 100644 --- a/ios/ClawChat/Tests/ClawChatTests/ConversationTests.swift +++ b/ios/ClawChat/Tests/ClawChatTests/ConversationTests.swift @@ -42,10 +42,15 @@ final class ConversationTests: XCTestCase { XCTAssertEqual(msg.attachments.first?.bytes.count, 3) } - func testSendAttachmentWhenNotReadyIsIgnored() { + func testSendAttachmentWhenNotReadyQueuesAsSending() { + // Same contract as send(): with no ready connection the message is + // queued in .sending for drainOutbox to pick up, not discarded. let c = Conversation() c.sendAttachment(kind: .image, bytes: Data([1, 2, 3]), filename: "x.png", mime: "image/png") - XCTAssertTrue(c.messages.isEmpty) + XCTAssertEqual(c.messages.count, 1) + XCTAssertEqual(c.messages.first?.delivery, .sending) + XCTAssertEqual(c.messages.first?.sender, .me) + XCTAssertEqual(c.messages.first?.attachments.first?.filename, "x.png") } func testIncomingAttachmentKindMapping() { diff --git a/plugin/src/allowlist.ts b/plugin/src/allowlist.ts index f9682bb..0196742 100644 --- a/plugin/src/allowlist.ts +++ b/plugin/src/allowlist.ts @@ -2,7 +2,7 @@ // and the agent. The Pilot trust handshake already proved the sender is who // they claim (X25519 + Ed25519); we just enforce *who is allowed at all*. -import { pilotAddrBase } from "./config.js"; +import { canonicalPilotAddr, pilotAddrBase } from "./config.js"; export type AllowlistDecision = | { allowed: true; peer: string } @@ -15,9 +15,19 @@ export function decideAllowlist( if (!rawSrcAddr || typeof rawSrcAddr !== "string") { return { allowed: false, reason: "malformed-src" }; } - const peer = pilotAddrBase(rawSrcAddr); + // The address grammar accepts hex in either case, so compare on the + // canonical form. `resolveAccount` already canonicalizes the configured + // set; the fallback scan covers sets assembled by other callers. + const peer = canonicalPilotAddr(pilotAddrBase(rawSrcAddr)); if (!allowlist.has(peer)) { - return { allowed: false, reason: "not-in-allowlist" }; + let found = false; + for (const entry of allowlist) { + if (canonicalPilotAddr(entry) === peer) { + found = true; + break; + } + } + if (!found) return { allowed: false, reason: "not-in-allowlist" }; } return { allowed: true, peer }; } diff --git a/plugin/src/config.ts b/plugin/src/config.ts index e42674a..831128d 100644 --- a/plugin/src/config.ts +++ b/plugin/src/config.ts @@ -65,7 +65,7 @@ export function resolveAccount( accountId, enabled: raw.enabled !== false, socketPath: raw.socketPath ?? DEFAULT_SOCKET_PATH, - allowlist: new Set(raw.allowlist ?? []), + allowlist: new Set((raw.allowlist ?? []).map(canonicalPilotAddr)), appPort: raw.appPort ?? DEFAULT_APP_PORT, handshakeTrustAutoApprove: raw.handshakeTrustAutoApprove !== false, sharedSecret: hasSecret ? raw.sharedSecret : undefined, @@ -78,6 +78,17 @@ export function isValidPilotAddress(addr: string): boolean { return typeof addr === "string" && PILOT_ADDR_RE.test(addr); } +/** + * Put a pilot address in the single case the rest of the plugin uses. The + * network id is decimal and the node id is hex, so upper-casing the whole + * string only touches the hex groups. Matches the form `nodeIdToAddress` + * emits in peer-address.ts, so addresses derived from either path compare + * equal. + */ +export function canonicalPilotAddr(addr: string): string { + return addr.toUpperCase(); +} + /** Strip an optional `:PORT` suffix from a pilot address-with-port. */ export function pilotAddrBase(addrWithMaybePort: string): string { const lastColon = addrWithMaybePort.lastIndexOf(":"); diff --git a/plugin/src/inbound.ts b/plugin/src/inbound.ts index 52fc0e6..e05ad28 100644 --- a/plugin/src/inbound.ts +++ b/plugin/src/inbound.ts @@ -11,7 +11,7 @@ import { mkdirSync, writeFileSync } from "node:fs"; import { join } from "node:path"; import { tmpdir } from "node:os"; -import type { ResolvedPilotAccount } from "./config.js"; +import { canonicalPilotAddr, pilotAddrBase, type ResolvedPilotAccount } from "./config.js"; import { createAegisScan, type AegisScan } from "./aegis-scan.js"; import { decideAllowlist } from "./allowlist.js"; import type { PeerAddressCache } from "./peer-address.js"; @@ -83,6 +83,12 @@ export type InboundDeps = { * envelopes are dropped with a warning. Default 25 MiB. */ maxMediaBytes?: number; + /** + * How far an envelope's `ts` may sit from local time, in either direction, + * for the HMAC path to accept it. Only applies when authorization came + * from the shared secret rather than the allowlist. Default 2 minutes. + */ + hmacMaxSkewMs?: number; /** * If set, the pipeline sends an `ack` envelope back to the peer once a * message has been fully reassembled (text or media). Lets senders @@ -122,6 +128,7 @@ export class InboundPipeline { private recentOrder: Array<{ id: string; ts: number }> = []; private readonly mediaDir: string; private readonly maxMediaBytes: number; + private readonly hmacMaxSkewMs: number; private readonly aegisScan: AegisScan; constructor(deps: InboundDeps) { @@ -129,6 +136,7 @@ export class InboundPipeline { this.recent = deps.recentIds ?? new Set(); this.mediaDir = deps.mediaDir ?? join(tmpdir(), "claw-pilot-inbound"); this.maxMediaBytes = deps.maxMediaBytes ?? 25 * 1024 * 1024; + this.hmacMaxSkewMs = deps.hmacMaxSkewMs ?? 120_000; this.aegisScan = deps.aegisScan ?? createAegisScan(); try { mkdirSync(this.mediaDir, { recursive: true }); @@ -199,7 +207,22 @@ export class InboundPipeline { const hmacOK = this.deps.account.sharedSecret ? await verifyEnvelope(env, this.deps.account.sharedSecret) : false; - if (!hmacOK) { + // The bypass additionally requires the envelope's `ts` to be inside the + // accepted window. `ts` is covered by the HMAC, so the window bounds how + // long a given signed envelope stays usable. Outside it, authorization + // falls back to the allowlist instead of being granted by the secret. + const skewMs = Math.abs(Date.now() - env.ts); + const hmacFresh = hmacOK && Number.isFinite(env.ts) && skewMs <= this.hmacMaxSkewMs; + if (hmacOK && !hmacFresh) { + this.deps.logger.warn("pilot inbound: HMAC envelope timestamp outside window", { + srcAddr: dg.srcAddr, + id: env.id, + ts: env.ts, + skewMs: Number.isFinite(skewMs) ? skewMs : null, + maxSkewMs: this.hmacMaxSkewMs, + }); + } + if (!hmacFresh) { const decision = decideAllowlist(dg.srcAddr, this.deps.account.allowlist); if (!decision.allowed) { this.deps.logger.warn("pilot inbound: dropped — not allowed", { @@ -211,11 +234,7 @@ export class InboundPipeline { } peer = decision.peer; } else { - // Strip any port suffix for consistent logging / dispatch. - const colonIdx = peer.lastIndexOf(":"); - if (colonIdx > 0 && /^\d+$/.test(peer.slice(colonIdx + 1))) { - peer = peer.slice(0, colonIdx); - } + peer = canonicalPilotAddr(pilotAddrBase(peer)); this.deps.logger.debug?.("pilot inbound: HMAC verified — bypassing allowlist", { srcAddr: peer, id: env.id, diff --git a/plugin/src/wire.ts b/plugin/src/wire.ts index 66d4e5e..cca2c41 100644 --- a/plugin/src/wire.ts +++ b/plugin/src/wire.ts @@ -381,12 +381,21 @@ export type ReassembledMedia = { caption?: string; }; +/** + * Ceiling on the summed payload of all chunks held for one in-flight media + * envelope. Chosen above the inbound pipeline's own 25 MiB attachment cap so + * that cap stays the one that governs completed transfers. + */ +export const MAX_MEDIA_REASSEMBLY_BYTES = 32 * 1024 * 1024; + /** Reassembler for media envelopes — collects binary chunks into a Buffer. */ export class MediaReassembler { private parts = new Map< string, { received: Map; + /** Decoded payload bytes currently held, summed across `received`. */ + bytes: number; total: number; firstTs: number; header?: { @@ -400,6 +409,12 @@ export class MediaReassembler { } >(); + private readonly maxBytes: number; + + constructor(maxBytes: number = MAX_MEDIA_REASSEMBLY_BYTES) { + this.maxBytes = maxBytes; + } + gc(now = Date.now(), maxAgeMs = 60_000): void { for (const [id, st] of this.parts) { if (now - st.firstTs > maxAgeMs) this.parts.delete(id); @@ -411,12 +426,24 @@ export class MediaReassembler { if (seq < 1 || seq > total) return null; let st = this.parts.get(id); if (!st) { - st = { received: new Map(), total, firstTs: env.ts }; + st = { received: new Map(), bytes: 0, total, firstTs: env.ts }; this.parts.set(id, st); } if (st.total !== total) return null; + // Re-sending a seq replaces the held chunk, so swap its contribution + // rather than adding to the running total. + const incoming = Buffer.byteLength(env.data, "base64"); + const previous = st.received.get(seq); + const projected = + st.bytes - (previous ? Buffer.byteLength(previous.data, "base64") : 0) + incoming; + if (projected > this.maxBytes) { + this.parts.delete(id); + return null; + } + st.received.set(seq, env); + st.bytes = projected; if (seq === 1) { st.header = { from: env.from, @@ -457,13 +484,32 @@ export class MediaReassembler { } } +/** + * Ceiling on the summed text of all chunks held for one in-flight message. + * A single envelope carries at most MAX_ENVELOPE_BYTES, so this still allows + * text far longer than any interactive message. + */ +export const MAX_TEXT_REASSEMBLY_BYTES = 1024 * 1024; + /** State for reassembling chunked messages keyed by envelope id. */ export class Reassembler { private parts = new Map< string, - { received: Map; total: number; firstTs: number } + { + received: Map; + /** Text bytes currently held, summed across `received`. */ + bytes: number; + total: number; + firstTs: number; + } >(); + private readonly maxBytes: number; + + constructor(maxBytes: number = MAX_TEXT_REASSEMBLY_BYTES) { + this.maxBytes = maxBytes; + } + /** Drop reassembly state for ids older than `maxAgeMs`. */ gc(now: number = Date.now(), maxAgeMs = 60_000): void { for (const [id, state] of this.parts) { @@ -484,11 +530,24 @@ export class Reassembler { if (seq === undefined || seq < 1 || seq > total) return null; let st = this.parts.get(env.id); if (!st) { - st = { received: new Map(), total, firstTs: env.ts }; + st = { received: new Map(), bytes: 0, total, firstTs: env.ts }; this.parts.set(env.id, st); } if (st.total !== total) return null; // contradictory total → drop + + // Re-sending a seq replaces the held chunk, so swap its contribution + // rather than adding to the running total. + const incoming = Buffer.byteLength(env.text ?? "", "utf8"); + const previous = st.received.get(seq); + const projected = + st.bytes - (previous ? Buffer.byteLength(previous.text ?? "", "utf8") : 0) + incoming; + if (projected > this.maxBytes) { + this.parts.delete(env.id); + return null; + } + st.received.set(seq, env); + st.bytes = projected; if (st.received.size < total) return null; let combined = ""; for (let i = 1; i <= total; i++) { diff --git a/plugin/tests/allowlist.test.ts b/plugin/tests/allowlist.test.ts index cdd920e..75c2c9c 100644 --- a/plugin/tests/allowlist.test.ts +++ b/plugin/tests/allowlist.test.ts @@ -38,10 +38,39 @@ describe("decideAllowlist", () => { }); }); - it("does not silently accept on a typo (case-sensitive net id)", () => { - // The pattern is hex case-insensitive for the hex parts, but the *exact* - // form we store should match the configured form. Strict equality here: - expect(decideAllowlist("1:0000.0000.aaaa", allow).allowed).toBe(false); + it("matches an allowlisted address whose hex differs only in case", () => { + // isValidPilotAddress accepts hex in either case, so both spellings name + // the same peer and both must resolve to the canonical form. + expect(decideAllowlist("1:0000.0000.aaaa", allow)).toEqual({ + allowed: true, + peer: "1:0000.0000.AAAA", + }); + expect(decideAllowlist("2:1234.5678.9abc:7777", allow)).toEqual({ + allowed: true, + peer: "2:1234.5678.9ABC", + }); + }); + + it("still rejects an address that is genuinely absent", () => { + expect(decideAllowlist("1:0000.0000.aaab", allow).allowed).toBe(false); + expect(decideAllowlist("9:0000.0000.aaaa", allow).allowed).toBe(false); + }); + + it("matches when the configured set itself is lower-case", () => { + const lower = new Set(["1:0000.0000.aaaa"]); + expect(decideAllowlist("1:0000.0000.AAAA", lower)).toEqual({ + allowed: true, + peer: "1:0000.0000.AAAA", + }); + }); +}); + +describe("resolveAccount — allowlist canonicalization", () => { + it("stores allowlist entries in canonical case", () => { + const acc = resolveAccount({ allowlist: ["1:00ab.cdef.0000"] }); + expect(acc.allowlist.has("1:00AB.CDEF.0000")).toBe(true); + expect(decideAllowlist("1:00ab.cdef.0000", acc.allowlist).allowed).toBe(true); + expect(decideAllowlist("1:00AB.CDEF.0000", acc.allowlist).allowed).toBe(true); }); }); diff --git a/plugin/tests/hmac.test.ts b/plugin/tests/hmac.test.ts index a257617..5f1259c 100644 --- a/plugin/tests/hmac.test.ts +++ b/plugin/tests/hmac.test.ts @@ -274,3 +274,99 @@ describe("resolveAccount — sharedSecret rules", () => { expect((back as { hmac?: string }).hmac).toBe("AAAA"); }); }); + +describe("InboundPipeline — HMAC envelope timestamp window", () => { + async function signedDatagram(ts: number, id = "t1") { + const env = { + v: WIRE_VERSION, + kind: "user" as const, + id, + ts, + text: "hello", + }; + const hmac = await signEnvelope(env, SECRET); + return encodeEnvelope({ ...env, hmac }); + } + + function makePipeline(hmacMaxSkewMs?: number) { + const dispatched: InboundDispatchInput[] = []; + const logger = silentLogger(); + const transport = new FakeTransport(); + const account = resolveAccount({ allowlist: [ALICE], sharedSecret: SECRET }); + const pipeline = new InboundPipeline({ + account, + dispatch: async (m) => { + dispatched.push(m); + }, + logger, + aegisScan: async () => ({ blocked: false, rule: "" }), + ...(hmacMaxSkewMs === undefined ? {} : { hmacMaxSkewMs }), + }); + pipeline.attach(transport); + return { dispatched, logger, transport, pipeline }; + } + + it("accepts an HMAC envelope whose ts is inside the window", async () => { + const { dispatched, transport, pipeline } = makePipeline(60_000); + transport.emitDatagram({ + srcAddr: STRANGER, + srcPort: 0, + dstPort: 7777, + data: await signedDatagram(Date.now() - 5_000), + }); + await new Promise((r) => setImmediate(r)); + expect(dispatched).toHaveLength(1); + pipeline.stop(); + }); + + it("drops an HMAC envelope whose ts is older than the window", async () => { + const { dispatched, logger, transport, pipeline } = makePipeline(60_000); + transport.emitDatagram({ + srcAddr: STRANGER, + srcPort: 0, + dstPort: 7777, + data: await signedDatagram(Date.now() - 10 * 60_000), + }); + await new Promise((r) => setImmediate(r)); + // Outside the window the secret grants nothing, so STRANGER is left to + // the allowlist — which does not contain it. + expect(dispatched).toHaveLength(0); + expect(logger.warn).toHaveBeenCalledWith( + "pilot inbound: HMAC envelope timestamp outside window", + expect.objectContaining({ srcAddr: STRANGER }), + ); + expect(logger.warn).toHaveBeenCalledWith( + "pilot inbound: dropped — not allowed", + expect.objectContaining({ srcAddr: STRANGER }), + ); + pipeline.stop(); + }); + + it("drops an HMAC envelope whose ts is far in the future", async () => { + const { dispatched, transport, pipeline } = makePipeline(60_000); + transport.emitDatagram({ + srcAddr: STRANGER, + srcPort: 0, + dstPort: 7777, + data: await signedDatagram(Date.now() + 10 * 60_000), + }); + await new Promise((r) => setImmediate(r)); + expect(dispatched).toHaveLength(0); + pipeline.stop(); + }); + + it("a stale envelope from an allowlisted peer is still delivered", async () => { + // The window gates the secret-based bypass only. An address that is on + // the allowlist is authorized on that basis, as before. + const { dispatched, transport, pipeline } = makePipeline(60_000); + transport.emitDatagram({ + srcAddr: ALICE, + srcPort: 0, + dstPort: 7777, + data: await signedDatagram(Date.now() - 10 * 60_000, "t2"), + }); + await new Promise((r) => setImmediate(r)); + expect(dispatched).toHaveLength(1); + pipeline.stop(); + }); +}); diff --git a/plugin/tests/wire-coverage.test.ts b/plugin/tests/wire-coverage.test.ts index f0f340a..09838c3 100644 --- a/plugin/tests/wire-coverage.test.ts +++ b/plugin/tests/wire-coverage.test.ts @@ -107,3 +107,76 @@ describe("verifyEnvelope — defensive paths", () => { expect(await verifyEnvelope(signed, SECRET)).toBe(true); }); }); + +describe("Reassembler byte cap", () => { + function textPart(id: string, seq: number, total: number, text: string): UserMessage { + return { v: WIRE_VERSION, kind: "user", id, ts: 1_000, text, seq, total }; + } + + it("drops the in-flight message once held text passes the cap", () => { + const r = new Reassembler(100); + const id = newId(); + // Two 40-byte chunks fit; the third pushes the total to 120 > 100. + expect(r.push(textPart(id, 1, 5, "a".repeat(40)))).toBeNull(); + expect(r.push(textPart(id, 2, 5, "b".repeat(40)))).toBeNull(); + expect(r.push(textPart(id, 3, 5, "c".repeat(40)))).toBeNull(); + // State was discarded, so the remaining chunks can never complete it. + expect(r.push(textPart(id, 4, 5, "d"))).toBeNull(); + expect(r.push(textPart(id, 5, 5, "e"))).toBeNull(); + }); + + it("still assembles a message that stays under the cap", () => { + const r = new Reassembler(100); + const id = newId(); + expect(r.push(textPart(id, 1, 2, "hello "))).toBeNull(); + const out = r.push(textPart(id, 2, 2, "world")); + expect(out?.text).toBe("hello world"); + }); + + it("does not let a re-sent seq inflate the running total", () => { + const r = new Reassembler(100); + const id = newId(); + // The same 40-byte chunk replayed many times replaces itself each time, + // so the held total stays at 40 and assembly still succeeds. + for (let i = 0; i < 20; i++) { + expect(r.push(textPart(id, 1, 2, "a".repeat(40)))).toBeNull(); + } + const out = r.push(textPart(id, 2, 2, "b".repeat(40))); + expect(out?.text).toBe("a".repeat(40) + "b".repeat(40)); + }); +}); + +describe("MediaReassembler byte cap", () => { + function mediaPart(id: string, seq: number, total: number, bytes: number): MediaMessage { + return { + v: WIRE_VERSION, + kind: "media", + from: "user", + media: "file", + id, + ts: 1_000, + data: Buffer.alloc(bytes, 7).toString("base64"), + seq, + total, + ...(seq === 1 ? { filename: "x.bin", totalBytes: bytes * total } : {}), + }; + } + + it("drops the in-flight media once held payload passes the cap", () => { + const r = new MediaReassembler(100); + const id = newId(); + expect(r.push(mediaPart(id, 1, 5, 40))).toBeNull(); + expect(r.push(mediaPart(id, 2, 5, 40))).toBeNull(); + expect(r.push(mediaPart(id, 3, 5, 40))).toBeNull(); + expect(r.push(mediaPart(id, 4, 5, 40))).toBeNull(); + expect(r.push(mediaPart(id, 5, 5, 40))).toBeNull(); + }); + + it("still assembles media that stays under the cap", () => { + const r = new MediaReassembler(100); + const id = newId(); + expect(r.push(mediaPart(id, 1, 2, 40))).toBeNull(); + const out = r.push(mediaPart(id, 2, 2, 40)); + expect(out?.bytes.length).toBe(80); + }); +});