Cleaner Implementation

[apoco]

Thanks so much for sharing this library and for accepting my bug fix. I ended up creating a refactored version of your code while investigating the bug. If you're interested, here's what I put together. I think it utilizes Node streaming concepts a bit more cleanly. Wanted to share it just in case you're interested in these same ideas since you were generous enough to share your code.

import calculateEtag from "etag";
import { Context, Middleware } from "koa";
import { Stream, Transform, TransformCallback } from "stream";

type Options = {
  weak?: boolean;
  sizeLimit?: number;
};

const DEFAULT_SIZE_LIMIT = 1024 * 1024;

export default function etag(options?: Options): Middleware {
  return async (ctx, next) => {
    await next();

    const body = ctx.body;
    if (!body || ctx.response.get("etag") || ctx.status / 100 !== 2) {
      return;
    }

    const outputStream = new ETagStream({ ctx, ...options });
    if (body instanceof Stream) {
      body.pipe(outputStream);
    } else {
      outputStream.end(
        typeof body === "string" || Buffer.isBuffer(body)
          ? body
          : JSON.stringify(body)
      );
    }

    ctx.body = outputStream;
  };
}

class ETagStream extends Transform {
  private ctx: Context;
  private sizeLimit: number;
  private weak: boolean;

  private isCalculating: boolean = true;
  private contentSize: number = 0;
  private bufferedChunks: Array<Buffer> = [];

  constructor({
    ctx,
    sizeLimit = DEFAULT_SIZE_LIMIT,
    weak = false,
  }: {
    ctx: Context;
    sizeLimit?: number;
    weak?: boolean;
  }) {
    super({ writableHighWaterMark: sizeLimit });

    this.ctx = ctx;
    this.sizeLimit = sizeLimit;
    this.weak = weak;
  }

  _transform(
    chunk: Buffer,
    _encoding: BufferEncoding,
    callback: TransformCallback
  ) {
    if (!this.isCalculating) {
      // We've already cancelled calculating the ETag, so just pass through
      return void callback(null, chunk);
    }

    this.bufferedChunks.push(chunk);
    this.contentSize += chunk.length;
    if (this.contentSize > this.sizeLimit) {
      // Bail on completing the calculation and flush the buffered data so far
      this.isCalculating = false;
      return void this._flush(callback);
    }

    // Hold the data until we have enough data to calculate the ETag
    return void callback();
  }

  _flush(callback: TransformCallback) {
    if (this.isCalculating) {
      // We've reached the end of the stream; finalize etag calculation
      const entity = Buffer.concat(this.bufferedChunks);
      this.ctx.response.etag = calculateEtag(entity, { weak: this.weak });
    }

    // Flush the buffered data
    let bufferedChunk;
    while ((bufferedChunk = this.bufferedChunks.shift())) {
      this.push(bufferedChunk);
    }

    return void callback();
  }
}

[masx200]

thank you

[masx200]

koa-stream-etag/index.mjs

Line 52 in 8f8b651

const [stream1, stream2] = ReadableStream.tee();

ReadableStream.tee();

I'm ready to ReFactor it, but it's not available for the time being because of nodejs's bug.

[masx200]

nodejs/node#43325