feat(etag): add getBody option

src/middleware/etag/digest.ts

@@ -9,23 +9,25 @@ const mergeBuffers = (buffer1: ArrayBuffer | undefined, buffer2: Uint8Array): Ui
 }
 
 export const generateDigest = async (
-  stream: ReadableStream<Uint8Array> | null,
+  input: ReadableStream<Uint8Array> | ArrayBuffer | null,
   generator: (body: Uint8Array) => ArrayBuffer | Promise<ArrayBuffer>
 ): Promise<string | null> => {
-  if (!stream) {
+  if (!input) {
     return null
   }
 
   let result: ArrayBuffer | undefined = undefined
-
-  const reader = stream.getReader()
-  for (;;) {
-    const { value, done } = await reader.read()
-    if (done) {
-      break
+  if (input instanceof ArrayBuffer) {
+    result = await generator(new Uint8Array(input))
+  } else {
+    const reader = input.getReader()
+    for (;;) {
+      const { value, done } = await reader.read()
+      if (done) {
+        break
+      }
+      result = await generator(mergeBuffers(result, value))
     }
-
-    result = await generator(mergeBuffers(result, value))
   }
 
   if (!result) {

src/middleware/etag/index.test.ts

@@ -290,6 +290,57 @@ describe('Etag Middleware', () => {
     expect(res.status).toBe(304)
   })
 
+  it('Should use custom getBody function', async () => {
+    const customBuffer = new ArrayBuffer(1)
+
+    const app = new Hono()
+    app.use(
+      '/etag/*',
+      etag({
+        getBody: async () => customBuffer,
+      })
+    )
+    app.get('/etag/custom', (c) => {
+      return c.text('Original response content')
+    })
+
+    const res = await app.request('http://localhost/etag/custom')
+    expect(res.headers.get('ETag')).not.toBeFalsy()
+    // ETag should be calculated from custom body, not the actual response
+    expect(res.headers.get('ETag')).toBe('"5ba93c9db0cff93f52b521d7420e43f6eda2784f"')
+    expect(await res.text()).toBe('Original response content')
+  })
+
+  it('Should work with getBody and custom generateDigest', async () => {
+    const customBuffer = new ArrayBuffer(1)
+
+    const app = new Hono()
+    app.use(
+      '/etag/*',
+      etag({
+        getBody: async () => customBuffer,
+        generateDigest: (body) =>
+          crypto.subtle.digest(
+            {
+              name: 'SHA-256',
+            },
+            body
+          ),
+      })
+    )
+    app.get('/etag/custom-both', (c) => {
+      return c.text('Different content')
+    })
+
+    const res = await app.request('http://localhost/etag/custom-both')
+    expect(res.headers.get('ETag')).not.toBeFalsy()
+    // Should use SHA-256 hash of the custom buffer
+    expect(res.headers.get('ETag')).toBe(
+      '"6e340b9cffb37a989ca544e6bb780a2c78901d3fb33738768511a30617afa01d"'
+    )
+    expect(await res.text()).toBe('Different content')
+  })
+
   describe('When crypto is not available', () => {
     let _crypto: Crypto | undefined
     beforeAll(() => {

src/middleware/etag/index.ts

@@ -3,13 +3,15 @@
  * ETag Middleware for Hono.
  */
 
+import type { Context } from '../../context'
 import type { MiddlewareHandler } from '../../types'
 import { generateDigest } from './digest'
 
 type ETagOptions = {
   retainedHeaders?: string[]
   weak?: boolean
   generateDigest?: (body: Uint8Array) => ArrayBuffer | Promise<ArrayBuffer>
+  getBody?: (c: Context) => ArrayBuffer | Promise<ArrayBuffer>
 }
 
 /**
@@ -64,6 +66,9 @@ function initializeGenerator(
  * @param {function(Uint8Array): ArrayBuffer | Promise<ArrayBuffer>} [options.generateDigest] -
  * A custom digest generation function. By default, it uses 'SHA-1'
  * This function is called with the response body as a `Uint8Array` and should return a hash as an `ArrayBuffer` or a Promise of one.
+ * @param {function(Context): ArrayBuffer | Promise<ArrayBuffer>} [options.getBody] -
+ * A custom function to get the body for ETag calculation. By default, it uses the response body.
+ * Useful for customizing what data is used to generate the ETag or in environments like AWS Lambda where `c.res.clone()` is not supported.
  * @returns {MiddlewareHandler} The middleware handler function.
  *
  * @example
@@ -93,7 +98,10 @@ export const etag = (options?: ETagOptions): MiddlewareHandler => {
       if (!generator) {
         return
       }
-      const hash = await generateDigest(res.clone().body, generator)
+
+      const body = options?.getBody ? await options.getBody(c) : res.clone().body
+      const hash = await generateDigest(body, generator)
+
       if (hash === null) {
         return
       }