Honour the zstd dictionary option in node:zlib

src/node/internal/internal_zlib.ts

@@ -16,6 +16,7 @@ import {
   Zlib,
   Brotli,
   Zstd,
+  normalizeZstdOptions,
   zstdInitCParamsArray,
   zstdInitDParamsArray,
   kMaxZstdCParam,
@@ -345,26 +346,28 @@ export function zstdDecompressSync(
   data: ArrayBufferView | string,
   options: ZstdOptions = {}
 ): ZlibResult {
-  if (!options.info) {
+  const opts = normalizeZstdOptions(options);
+  if (!opts.info) {
     // Fast path, where we send the data directly to C++
-    return Buffer.from(zlibUtil.zstdDecompressSync(data, options));
+    return Buffer.from(zlibUtil.zstdDecompressSync(data, opts));
   }
 
   // Else, use the Engine class in sync mode
-  return processChunk(new ZstdDecompress(options), data);
+  return processChunk(new ZstdDecompress(opts), data);
 }
 
 export function zstdCompressSync(
   data: ArrayBufferView | string,
   options: ZstdOptions = {}
 ): ZlibResult {
-  if (!options.info) {
+  const opts = normalizeZstdOptions(options);
+  if (!opts.info) {
     // Fast path, where we send the data directly to C++
-    return Buffer.from(zlibUtil.zstdCompressSync(data, options));
+    return Buffer.from(zlibUtil.zstdCompressSync(data, opts));
   }
 
   // Else, use the Engine class in sync mode
-  return processChunk(new ZstdCompress(options), data);
+  return processChunk(new ZstdCompress(opts), data);
 }
 
 export function zstdDecompress(
@@ -376,10 +379,11 @@ export function zstdDecompress(
     optionsOrCallback,
     callbackOrUndefined
   );
+  const opts = normalizeZstdOptions(options);
 
-  if (!options.info) {
+  if (!opts.info) {
     // Fast path
-    zlibUtil.zstdDecompress(data, options, (res) => {
+    zlibUtil.zstdDecompress(data, opts, (res) => {
       queueMicrotask(() => {
         if (res instanceof Error) {
           callback(res);
@@ -392,7 +396,7 @@ export function zstdDecompress(
     return;
   }
 
-  processChunkCaptureError(new ZstdDecompress(options), data, callback);
+  processChunkCaptureError(new ZstdDecompress(opts), data, callback);
 }
 
 export function zstdCompress(
@@ -404,10 +408,11 @@ export function zstdCompress(
     optionsOrCallback,
     callbackOrUndefined
   );
+  const opts = normalizeZstdOptions(options);
 
-  if (!options.info) {
+  if (!opts.info) {
     // Fast path
-    zlibUtil.zstdCompress(data, options, (res) => {
+    zlibUtil.zstdCompress(data, opts, (res) => {
       queueMicrotask(() => {
         if (res instanceof Error) {
           callback(res);
@@ -420,7 +425,7 @@ export function zstdCompress(
     return;
   }
 
-  processChunkCaptureError(new ZstdCompress(options), data, callback);
+  processChunkCaptureError(new ZstdCompress(opts), data, callback);
 }
 export class Gzip extends Zlib {
   constructor(options: ZlibOptions) {

src/node/internal/internal_zlib_base.ts

@@ -826,6 +826,35 @@ export const kMaxZstdDParam = Math.max(
 );
 export const zstdInitDParamsArray = new Int32Array(kMaxZstdDParam + 1);
 
+// Node accepts a zstd `dictionary` as an ArrayBufferView or an ArrayBuffer, and throws
+// ERR_INVALID_ARG_TYPE for any other type, null included (lib/zlib.js, class Zstd, since
+// nodejs/node#65867; earlier releases ignored it). Normalize here so that both the stream
+// path and the convenience functions' fast path behave the same way.
+export function normalizeZstdDictionary(
+  dictionary: ZstdOptions['dictionary']
+): ArrayBufferView | undefined {
+  if (dictionary === undefined || isArrayBufferView(dictionary)) {
+    return dictionary;
+  }
+  if (isAnyArrayBuffer(dictionary)) {
+    return new Uint8Array(dictionary);
+  }
+  throw new ERR_INVALID_ARG_TYPE(
+    'options.dictionary',
+    ['Buffer', 'TypedArray', 'DataView', 'ArrayBuffer'],
+    dictionary
+  );
+}
+
+// Returns `options` unchanged unless its dictionary needed normalizing, so the common case
+// allocates nothing.
+export function normalizeZstdOptions(options: ZstdOptions): ZstdOptions {
+  const dictionary = normalizeZstdDictionary(options.dictionary);
+  return dictionary === options.dictionary
+    ? options
+    : { ...options, dictionary };
+}
+
 const zstdDefaultOptions: ZlibDefaultOptions = {
   flush: CONST_ZSTD_e_continue,
   finishFlush: CONST_ZSTD_e_end,
@@ -883,7 +912,8 @@ export class Zstd extends ZlibBase {
         () => {
           queueMicrotask(processCallback.bind(handle));
         },
-        pledgedSrcSize
+        pledgedSrcSize,
+        normalizeZstdDictionary(options?.dictionary)
       )
     ) {
       throw new ERR_ZLIB_INITIALIZATION_FAILED();

src/node/internal/zlib.d.ts

@@ -291,6 +291,10 @@ export interface ZstdOptions {
     | undefined;
   maxOutputLength?: number | undefined;
   pledgedSrcSize?: number | undefined;
+  // Declared as a view, like ZlibOptions above, though an ArrayBuffer is also accepted at
+  // runtime and any other type throws. See normalizeZstdDictionary() in
+  // internal_zlib_base.ts.
+  dictionary?: ArrayBufferView | undefined;
   // Not specified in NodeJS docs but the tests expect it
   info?: boolean | undefined;
 }
@@ -370,7 +374,8 @@ export class ZstdDecoder extends CompressionStream {
     params: Int32Array,
     writeResult: Uint32Array,
     writeCallback: () => void,
-    pledgedSrcSize?: number
+    pledgedSrcSize?: number,
+    dictionary?: ArrayBufferView
   ): boolean;
   params(): void;
 }
@@ -380,7 +385,8 @@ export class ZstdEncoder extends CompressionStream {
     params: Int32Array,
     writeResult: Uint32Array,
     writeCallback: () => void,
-    pledgedSrcSize?: number
+    pledgedSrcSize?: number,
+    dictionary?: ArrayBufferView
   ): boolean;
   params(): void;
 }

src/workerd/api/compression.c++

@@ -597,12 +597,29 @@ ZstdEncoderContext::ZstdEncoderContext(ZlibMode _mode)
     : ZstdContext(_mode),
       cctx_(kj::disposeWith<zstdFreeCCtx>(ZSTD_createCCtx())) {}
 
-kj::Maybe<CompressionError> ZstdEncoderContext::initialize(uint64_t pledgedSrcSize) {
+kj::Maybe<CompressionError> ZstdEncoderContext::initialize(
+    uint64_t pledgedSrcSize, kj::ArrayPtr<const kj::byte> dictionary) {
   if (cctx_.get() == nullptr) {
     return CompressionError(
         "Could not initialize Zstd instance"_kj, "ERR_ZLIB_INITIALIZATION_FAILED"_kj, -1);
   }
 
+  if (dictionary.size() > 0) {
+    // ZSTD_CCtx_loadDictionary() copies the dictionary into the context, so `dictionary` does
+    // not need to outlive this call. The content type is auto-detected: a buffer starting with
+    // the zstd dictionary magic is read as a trained dictionary, anything else as raw content.
+    // Loading is deferred until the first frame begins, so the parameters set by setParams()
+    // afterwards still apply to the dictionary's tables. It also means a malformed trained
+    // dictionary is not detected here: it fails in work() when the first frame begins, which
+    // is where Node reports it too. On a fresh context this call can only fail to allocate.
+    size_t result = ZSTD_CCtx_loadDictionary(cctx_.get(), dictionary.begin(), dictionary.size());
+    if (ZSTD_isError(result)) {
+      error_ = ZSTD_getErrorCode(result);
+      return CompressionError(
+          "Failed to load zstd dictionary"_kj, "ERR_ZLIB_DICTIONARY_LOAD_FAILED"_kj, -1);
+    }
+  }
+
   if (pledgedSrcSize != ZSTD_CONTENTSIZE_UNKNOWN) {
     size_t result = ZSTD_CCtx_setPledgedSrcSize(cctx_.get(), pledgedSrcSize);
     KJ_IF_SOME(err, zstdCheckError(result, error_, "ERR_ZSTD_COMPRESSION_FAILED"_kj)) {
@@ -668,14 +685,34 @@ ZstdDecoderContext::ZstdDecoderContext(ZlibMode _mode)
     : ZstdContext(_mode),
       dctx_(kj::disposeWith<zstdFreeDCtx>(ZSTD_createDCtx())) {}
 
-kj::Maybe<CompressionError> ZstdDecoderContext::initialize() {
+kj::Maybe<CompressionError> ZstdDecoderContext::initialize(
+    kj::ArrayPtr<const kj::byte> dictionary) {
   // dctx_ is created in the constructor. It can only be nullptr if ZSTD_createDCtx()
   // failed due to memory allocation failure.
   if (dctx_.get() == nullptr) {
     return CompressionError(
         "Could not initialize Zstd instance"_kj, "ERR_ZLIB_INITIALIZATION_FAILED"_kj, -1);
   }
 
+  if (dictionary.size() > 0) {
+    // As with the encoder, the bytes are copied into the context and the content type is
+    // auto-detected. Note that a raw-content dictionary carries no dictionary ID, so reading a
+    // frame written against a different one is not rejected as ZSTD_error_dictionary_wrong: it
+    // fails as corrupt, or decodes to different bytes if the frame carries no checksum. That is
+    // zstd's behaviour and matches what Node does with the same calls.
+    //
+    // Unlike the encoder, the decoder parses a trained dictionary's entropy tables here, so a
+    // malformed one fails now. zstd reports that as ZSTD_error_memory_allocation, because the
+    // DDict it tried to build came back null, so the message leaves zstd's error name out and
+    // uses Node's exact text instead.
+    size_t result = ZSTD_DCtx_loadDictionary(dctx_.get(), dictionary.begin(), dictionary.size());
+    if (ZSTD_isError(result)) {
+      error_ = ZSTD_getErrorCode(result);
+      return CompressionError(
+          "Failed to load zstd dictionary"_kj, "ERR_ZLIB_DICTIONARY_LOAD_FAILED"_kj, -1);
+    }
+  }
+
   return kj::none;
 }
 

src/workerd/api/compression.h

@@ -461,7 +461,8 @@ class ZstdContext {
     jsg::Optional<jsg::Dict<int>> params;
     jsg::Optional<kj::uint> maxOutputLength;
     jsg::Optional<uint64_t> pledgedSrcSize;
-    JSG_STRUCT(flush, finishFlush, chunkSize, params, maxOutputLength, pledgedSrcSize);
+    jsg::Optional<kj::Array<kj::byte>> dictionary;
+    JSG_STRUCT(flush, finishFlush, chunkSize, params, maxOutputLength, pledgedSrcSize, dictionary);
   };
 
  protected:
@@ -480,7 +481,11 @@ class ZstdEncoderContext final: public ZstdContext {
   KJ_DISALLOW_COPY_AND_MOVE(ZstdEncoderContext);
 
   void work();
-  kj::Maybe<CompressionError> initialize(uint64_t pledgedSrcSize);
+  // An empty `dictionary` means no dictionary, matching Node.js, which passes an empty
+  // std::string_view for the same case. The bytes are copied into the context, so the
+  // caller does not need to keep them alive.
+  kj::Maybe<CompressionError> initialize(
+      uint64_t pledgedSrcSize, kj::ArrayPtr<const kj::byte> dictionary = nullptr);
   kj::Maybe<CompressionError> resetStream();
   kj::Maybe<CompressionError> setParams(int key, int value);
   kj::Maybe<CompressionError> getError() const;
@@ -501,7 +506,8 @@ class ZstdDecoderContext final: public ZstdContext {
   KJ_DISALLOW_COPY_AND_MOVE(ZstdDecoderContext);
 
   void work();
-  kj::Maybe<CompressionError> initialize();
+  // See the note on ZstdEncoderContext::initialize() regarding `dictionary`.
+  kj::Maybe<CompressionError> initialize(kj::ArrayPtr<const kj::byte> dictionary = nullptr);
   kj::Maybe<CompressionError> resetStream();
   kj::Maybe<CompressionError> setParams(int key, int value);
   kj::Maybe<CompressionError> getError() const;