Enforce a pledged zstd source size on one-shot compression and across reset()

src/workerd/api/compression.c++

@@ -608,6 +608,9 @@ kj::Maybe<CompressionError> ZstdEncoderContext::initialize(uint64_t pledgedSrcSi
     KJ_IF_SOME(err, zstdCheckError(result, error_, "ERR_ZSTD_COMPRESSION_FAILED"_kj)) {
       return kj::mv(err);
     }
+    // Start counting, so that work() can enforce the pledge even when zstd replaces it.
+    pledgedSrcSize_ = pledgedSrcSize;
+    consumedSrcSize_ = uint64_t(0);
   }
 
   return kj::none;
@@ -617,10 +620,26 @@ void ZstdEncoderContext::work() {
   JSG_REQUIRE(mode == ZlibMode::ZSTD_ENCODE, Error, "Mode should be ZSTD_ENCODE"_kj);
   JSG_REQUIRE(cctx_.get() != nullptr, Error, "Zstd context should not be null"_kj);
 
+  size_t inputPos = input_.pos;
   lastResult = ZSTD_compressStream2(cctx_.get(), &output_, &input_, flush_);
+  KJ_IF_SOME(consumed, consumedSrcSize_) {
+    consumed += input_.pos - inputPos;
+  }
 
   if (ZSTD_isError(lastResult)) {
     error_ = ZSTD_getErrorCode(lastResult);
+  } else if (lastResult == 0 && flush_ == ZSTD_e_end) {
+    // The frame is complete. If a size was pledged, compare it with what the frame actually
+    // consumed, since zstd may have replaced the pledge (see consumedSrcSize_ in the header).
+    // A frame that spanned several calls was already checked by zstd, which then fails in
+    // the branch above, so this only adds the check zstd skipped.
+    KJ_IF_SOME(consumed, consumedSrcSize_) {
+      if (consumed != pledgedSrcSize_) {
+        error_ = ZSTD_error_srcSize_wrong;
+      }
+    }
+    // Check once per frame. resetStream() starts the count again for the next one.
+    consumedSrcSize_ = kj::none;
   }
 }
 
@@ -630,6 +649,16 @@ kj::Maybe<CompressionError> ZstdEncoderContext::resetStream() {
     KJ_IF_SOME(err, zstdCheckError(result, error_, "ERR_ZSTD_COMPRESSION_FAILED"_kj)) {
       return kj::mv(err);
     }
+
+    // A session reset also sets the pledged size back to unknown, so pledge it again for the
+    // next frame and restart the count that enforces it.
+    if (pledgedSrcSize_ != ZSTD_CONTENTSIZE_UNKNOWN) {
+      result = ZSTD_CCtx_setPledgedSrcSize(cctx_.get(), pledgedSrcSize_);
+      KJ_IF_SOME(err, zstdCheckError(result, error_, "ERR_ZSTD_COMPRESSION_FAILED"_kj)) {
+        return kj::mv(err);
+      }
+      consumedSrcSize_ = uint64_t(0);
+    }
   }
   return kj::none;
 }

src/workerd/api/compression.h

@@ -480,6 +480,9 @@ class ZstdEncoderContext final: public ZstdContext {
   KJ_DISALLOW_COPY_AND_MOVE(ZstdEncoderContext);
 
   void work();
+  // Pledging a size other than ZSTD_CONTENTSIZE_UNKNOWN makes work() reject a frame whose
+  // input turns out to be a different length, and resetStream() carries the pledge over to
+  // the next frame.
   kj::Maybe<CompressionError> initialize(uint64_t pledgedSrcSize);
   kj::Maybe<CompressionError> resetStream();
   kj::Maybe<CompressionError> setParams(int key, int value);
@@ -490,6 +493,18 @@ class ZstdEncoderContext final: public ZstdContext {
   size_t lastResult = 0;
   kj::Own<ZSTD_CCtx> cctx_;
   ZSTD_ErrorCode error_ = ZSTD_error_no_error;
+
+  // The size passed to initialize(), kept so that resetStream() can pledge it again, since a
+  // session reset makes zstd forget it.
+  uint64_t pledgedSrcSize_ = ZSTD_CONTENTSIZE_UNKNOWN;
+
+  // Input consumed by the current frame, counted only while a size is pledged. zstd enforces
+  // a pledge by itself only when a frame spans more than one call: when the first call is
+  // also the last (ZSTD_e_end), it replaces the pledge with the real input size, so a wrong
+  // pledge would produce a valid frame and no error. The one-shot node:zlib functions always
+  // compress that way. work() compares this count with the pledge when the frame ends, which
+  // is what Node does too.
+  kj::Maybe<uint64_t> consumedSrcSize_;
 };
 
 class ZstdDecoderContext final: public ZstdContext {

src/workerd/api/node/tests/zlib-zstd-nodejs-test.js

@@ -446,3 +446,89 @@ export const zstdStreamLargeDecompressTest = {
     );
   },
 };
+
+// A pledged source size is a promise about how many bytes the frame will hold, and a wrong
+// one must be rejected. zstd only enforces it by itself when the input arrives over more
+// than one call: a first call that is also the last (ZSTD_e_end) makes zstd replace the
+// pledge with the real input size. The one-shot functions always compress that way, so the
+// encoder has to count the input itself, as Node does.
+const PLEDGE_INPUT = Buffer.from('pledged source size '.repeat(64));
+
+// Compresses through a stream and resolves to the output or the error, whichever comes.
+function compressThroughStream(stream, input) {
+  const { promise, resolve } = Promise.withResolvers();
+  const chunks = [];
+  stream.on('data', (chunk) => chunks.push(chunk));
+  stream.on('error', (err) => resolve({ err }));
+  stream.on('end', () => resolve({ out: Buffer.concat(chunks) }));
+  stream.end(input);
+  return promise;
+}
+
+export const zstdPledgedSrcSizeOneShotTest = {
+  async test() {
+    const input = PLEDGE_INPUT;
+
+    // A correct pledge still round-trips on every one-shot path.
+    for (const info of [false, true]) {
+      const result = zlib.zstdCompressSync(input, {
+        pledgedSrcSize: input.length,
+        info,
+      });
+      const compressed = info ? result.buffer : result;
+      assert(
+        zlib.zstdDecompressSync(compressed).equals(input),
+        'A correct pledge should round-trip'
+      );
+    }
+
+    // Too large and too small are both wrong. zstd's own override would hide either one.
+    for (const pledgedSrcSize of [input.length + 1, input.length - 1]) {
+      assert.throws(
+        () => zlib.zstdCompressSync(input, { pledgedSrcSize }),
+        /Src size is incorrect/,
+        `The fast path should reject a pledge of ${pledgedSrcSize}`
+      );
+      assert.throws(
+        () => zlib.zstdCompressSync(input, { pledgedSrcSize, info: true }),
+        /Src size is incorrect/,
+        `The engine path should reject a pledge of ${pledgedSrcSize}`
+      );
+
+      const { promise, resolve } = Promise.withResolvers();
+      zlib.zstdCompress(input, { pledgedSrcSize }, (err) => resolve(err));
+      const err = await promise;
+      assert.match(
+        err?.message ?? '',
+        /Src size is incorrect/,
+        `The async function should reject a pledge of ${pledgedSrcSize}`
+      );
+    }
+  },
+};
+
+// reset() starts a new frame, and the pledge has to apply to that frame too. zstd forgets
+// a pledge on a session reset, so the encoder has to set it again.
+export const zstdPledgedSrcSizeAfterResetTest = {
+  async test() {
+    const input = PLEDGE_INPUT;
+
+    const exact = zlib.createZstdCompress({ pledgedSrcSize: input.length });
+    exact.reset();
+    const good = await compressThroughStream(exact, input);
+    assert.ifError(good.err);
+    assert(
+      zlib.zstdDecompressSync(good.out).equals(input),
+      'A correct pledge should round-trip after reset()'
+    );
+
+    const wrong = zlib.createZstdCompress({ pledgedSrcSize: input.length + 1 });
+    wrong.reset();
+    const bad = await compressThroughStream(wrong, input);
+    assert.match(
+      bad.err?.message ?? '',
+      /Src size is incorrect/,
+      'A wrong pledge should still be rejected after reset()'
+    );
+  },
+};