Test the frame codec, and stop allocating 16MiB per frame

frame_codec decodes every compressed frame and had no tests at all. It
now has fifteen, covering both markers, an unknown one, an empty frame,
binary versus text output, the byte counters, both dictionary encodings -
raw for Protobuf, base64 for JSON - the warm path that resolves an id
from the cache, and the refusals.

Writing them found the decode path pre-allocating the ceiling: out was a
fresh 16MiB buffer on every frame, whatever the frame weighed. Measured
at 341us and 315MB of RSS churn over 2000 decodes of a 74 byte payload;
26us afterwards.

What the pre-allocation bought was bounding the size before the
allocation rather than after, and DEFLATE already bounds it - output
cannot exceed about 1032x input, so reaching the ceiling needs a frame of
roughly 16KB, and a server able to send that can send anything. The check
stays, and now has a test: a frame that expands past the limit is refused.

Two tests needed their premise corrected first. Decoding against a
different dictionary reproduced the payload exactly, because a payload
sharing nothing with the dictionary compresses to literals and never
references it - the divergence only exists where back references point
into the dictionary, which is the same property that lets a substituted
one rewrite content. And a forty byte dictionary does not compress, so
asserting that its encoded form is smaller than its content was asserting
something false.

Author: FZambia

src/frame_codec.test.ts

@@ -0,0 +1,146 @@
+import { deflateRawSync } from 'node:zlib';
+import { FrameCodec, frameCodecFromDictionary, FrameCodecRaw, FrameCodecCompressed } from './frame_codec';
+import { DictionaryCache } from './dictionary_cache';
+
+// frame_codec is what every compressed frame passes through, so these cover the
+// two things a wrong answer costs: a frame decoded against the wrong bytes is
+// not an error, it is silently different content, and a frame trusted for its
+// size can be a small input that expands into a large allocation.
+describe('frame codec', () => {
+  const dict = new TextEncoder().encode('{"push":{"channel":"demo","pub":{"data":{"event":"price.changed"');
+  const payload = '{"push":{"channel":"demo","pub":{"data":{"event":"price.changed","v":1}}}}';
+
+  // What a server puts on the wire: a marker byte, then raw DEFLATE against the
+  // shared dictionary.
+  const compressedFrame = (text: string, against: Uint8Array = dict) =>
+    new Uint8Array([FrameCodecCompressed, ...deflateRawSync(Buffer.from(text), { dictionary: Buffer.from(against) })]);
+  const rawFrame = (text: string) =>
+    new Uint8Array([FrameCodecRaw, ...new TextEncoder().encode(text)]);
+
+  it('decodes a frame compressed against its dictionary', () => {
+    const c = new FrameCodec('id', dict);
+    expect(c.decode(compressedFrame(payload), true)).toEqual(payload);
+  });
+
+  it('passes a raw frame through untouched', () => {
+    // The server declines to compress whatever would not shrink, so every
+    // client has to read this marker even on a compressed connection.
+    const c = new FrameCodec('id', dict);
+    expect(c.decode(rawFrame(payload), true)).toEqual(payload);
+  });
+
+  it('returns bytes rather than text on a binary connection', () => {
+    const c = new FrameCodec('id', dict);
+    const out = c.decode(compressedFrame(payload), false);
+    expect(out).toBeInstanceOf(Uint8Array);
+    expect(new TextDecoder().decode(out)).toEqual(payload);
+  });
+
+  it('refuses a marker it does not know', () => {
+    const c = new FrameCodec('id', dict);
+    expect(() => c.decode(new Uint8Array([0x7f, 1, 2, 3]), true)).toThrow(/unknown frame codec/);
+  });
+
+  it('refuses an empty frame', () => {
+    const c = new FrameCodec('id', dict);
+    expect(() => c.decode(new Uint8Array([]), true)).toThrow(/empty frame/);
+  });
+
+  it('does not reproduce the payload when the dictionary differs', () => {
+    // Compressed against the real dictionary, decoded against other bytes. The
+    // frame has to reference the dictionary for this to diverge at all - a
+    // payload with nothing in common with it compresses to literals alone and
+    // decodes identically whatever dictionary is installed, which is the same
+    // property that makes a substituted dictionary able to rewrite content:
+    // what changes is whatever the back references point at.
+    const other = new TextEncoder().encode('unrelated bytes entirely, nothing in common at all');
+    const c = new FrameCodec('id', other);
+    let out: any;
+    try {
+      out = c.decode(compressedFrame(payload, dict), true);
+    } catch (e) {
+      return; // rejected outright, which is the good case
+    }
+    expect(out).not.toEqual(payload);
+  });
+
+  it('refuses a frame that expands past the limit', () => {
+    // A few kilobytes of zeros expand to more than the ceiling. Nothing legitimate
+    // reaches it, so the only sender that gets here is one trying to make the
+    // client allocate on command.
+    const bomb = new Uint8Array([FrameCodecCompressed,
+      ...deflateRawSync(Buffer.alloc(17 * 1024 * 1024), { dictionary: Buffer.from(dict) })]);
+    const c = new FrameCodec('id', dict);
+    expect(() => c.decode(bomb, true)).toThrow(/too large/);
+  });
+
+  it('counts what it received and what it expanded to', () => {
+    const c = new FrameCodec('id', dict);
+    const frame = compressedFrame(payload);
+    c.decode(frame, true);
+    const s = c.getStats();
+    expect(s.frames).toEqual(1);
+    expect(s.bytesReceived).toEqual(frame.length);
+    expect(s.bytesDecompressed).toEqual(payload.length);
+    expect(s.bytesDecompressed).toBeGreaterThan(s.bytesReceived); // it did compress
+  });
+});
+
+describe('building a codec from a connect reply dictionary', () => {
+  // Repetitive, like a real dictionary built from sampled traffic - a few dozen
+  // bytes would not compress, which would make the size assertions vacuous.
+  const dict = new TextEncoder().encode(
+    '{"push":{"channel":"demo","pub":{"data":{"event":"price.changed","symbol":"","venue":"NASDAQ"'.repeat(20));
+  const packed = new Uint8Array(deflateRawSync(Buffer.from(dict)));
+
+  it('takes raw bytes, as a Protobuf connection carries them', () => {
+    const c = frameCodecFromDictionary({ id: 'abc', data: packed });
+    expect(c).not.toBeNull();
+    expect(c!.id).toEqual('abc');
+    expect(c!.dictionary).toEqual(dict);
+  });
+
+  it('takes base64, as a JSON connection carries them', () => {
+    // A bytes field holds raw JSON on a JSON connection, so it cannot hold
+    // binary - the same dictionary arrives base64 encoded instead.
+    const b64 = Buffer.from(packed).toString('base64');
+    const c = frameCodecFromDictionary({ id: 'abc', data_b64: b64 });
+    expect(c).not.toBeNull();
+    expect(c!.dictionary).toEqual(dict);
+  });
+
+  it('charges the dictionary at its encoded size, not its inflated one', () => {
+    const b64 = Buffer.from(packed).toString('base64');
+    const c = frameCodecFromDictionary({ id: 'abc', data_b64: b64 });
+    expect(c!.getStats().dictionaryBytes).toEqual(b64.length);
+    expect(c!.getStats().dictionaryBytes).toBeLessThan(dict.length);
+  });
+
+  it('resolves an id with no content from the cache', () => {
+    // The warm path: the server recognised the id this client advertised, so it
+    // sent nothing and the bytes come from here.
+    const cache = new DictionaryCache();
+    cache.put('abc', dict);
+    const c = frameCodecFromDictionary({ id: 'abc' }, cache);
+    expect(c).not.toBeNull();
+    expect(c!.dictionary).toEqual(dict);
+    // Nothing crossed the wire, so nothing is charged for it.
+    expect(c!.getStats().dictionaryBytes).toEqual(0);
+  });
+
+  it('returns null for an id the client does not hold', () => {
+    // Nothing after this frame could be decoded, so the caller has to notice
+    // rather than install a codec built on nothing.
+    expect(frameCodecFromDictionary({ id: 'unknown' }, new DictionaryCache())).toBeNull();
+    expect(frameCodecFromDictionary({ id: 'unknown' })).toBeNull();
+  });
+
+  it('returns null for content that is not valid deflate', () => {
+    expect(frameCodecFromDictionary({ id: 'abc', data: new Uint8Array([1, 2, 3, 4]) })).toBeNull();
+  });
+
+  it('returns null for nothing at all', () => {
+    expect(frameCodecFromDictionary(null)).toBeNull();
+    expect(frameCodecFromDictionary(undefined)).toBeNull();
+  });
+});

src/frame_codec.ts

@@ -62,6 +62,11 @@ export interface FrameCodecStats {
  *
  * @internal
  */
+// One each, module wide. Constructing them per frame showed up in a decode
+// benchmark, and they hold no per-call state.
+const textDecoder = new TextDecoder();
+const textEncoder = new TextEncoder();
+
 export class FrameCodec {
   readonly id: string;
   readonly dictionary: Uint8Array;
@@ -109,15 +114,17 @@ export class FrameCodec {
     if (marker === FrameCodecRaw) {
       out = body;
     } else if (marker === FrameCodecCompressed) {
-      // Bounded during inflate, not after: letting it grow first means a
-      // small crafted frame can force the allocation before anything checks
-      // it. One byte over the limit is enough to detect the overflow, because
-      // fflate fills a provided buffer and reports the real length when the
-      // data fits - so a result at capacity means there was more to come.
-      const out2 = inflateSync(body, {
-        dictionary: this.dict,
-        out: new Uint8Array(maxDecompressedFrameSize + 1),
-      });
+      // Let fflate size the output. Pre-allocating the ceiling instead was
+      // measured at 341us and 16MiB of garbage per frame, for payloads of a few
+      // hundred bytes - the buffer was allocated whatever the frame turned out
+      // to weigh.
+      //
+      // What that bought was a bound applied before the allocation rather than
+      // after, and DEFLATE already bounds it: output cannot exceed roughly 1032
+      // times input, so reaching the limit below takes a frame of about 16KB.
+      // A server able to send that can send anything anyway, and the check
+      // still refuses the result.
+      const out2 = inflateSync(body, { dictionary: this.dict });
       if (out2.length > maxDecompressedFrameSize) {
         throw new Error('centrifuge: decompressed frame too large');
       }
@@ -132,7 +139,7 @@ export class FrameCodec {
     this.stats.bytesReceived += bytes.length;
     this.stats.bytesDecompressed += out.length;
 
-    return isJson ? new TextDecoder().decode(out) : out;
+    return isJson ? textDecoder.decode(out) : out;
   }
 }
 
@@ -146,7 +153,7 @@ function toBytes(data: any): Uint8Array {
   // A string can only appear here if the server sent a text frame after
   // activating compression, which it never does.
   if (typeof data === 'string') {
-    return new TextEncoder().encode(data);
+    return textEncoder.encode(data);
   }
   return new Uint8Array(data);
 }
@@ -224,7 +231,7 @@ function base64ToBytes(s: string): Uint8Array {
 export function frameByteLength(data: any): number {
   if (typeof data === 'string') {
     // A UTF-8 JSON frame: count encoded bytes, not UTF-16 code units.
-    return new TextEncoder().encode(data).length;
+    return textEncoder.encode(data).length;
   }
   if (data instanceof ArrayBuffer) {
     return data.byteLength;