From e6d253d96fcbfb079ee6313d67d62c0fed84cb2d Mon Sep 17 00:00:00 2001
From: ryu <114303361+ryuapp@users.noreply.github.com>
Date: Sat, 13 Jul 2024 14:02:36 +0900
Subject: [PATCH] feat(middleware): introduce Request ID middleware (#3082)

* feat(middleware): introduce Request ID middleware

* fix not to accept empty string in header

* rename requestID to requestId

* pass the context to the generator option

* add typesVersions

* fix typo

Co-Authored-By: Taku Amano <taku@taaas.jp>

* change to generate id if validation fails

Co-Authored-By: Taku Amano <taku@taaas.jp>

* fix limit length test

---------

Co-authored-by: Taku Amano <taku@taaas.jp>
---
 jsr.json                                |   1 +
 package.json                            |   8 ++
 src/middleware/request-id/index.test.ts | 155 ++++++++++++++++++++++++
 src/middleware/request-id/index.ts      |   8 ++
 src/middleware/request-id/request-id.ts |  59 +++++++++
 5 files changed, 231 insertions(+)
 create mode 100644 src/middleware/request-id/index.test.ts
 create mode 100644 src/middleware/request-id/index.ts
 create mode 100644 src/middleware/request-id/request-id.ts

diff --git a/jsr.json b/jsr.json
index a4e3a65d..6b43f85f 100644
--- a/jsr.json
+++ b/jsr.json
@@ -49,6 +49,7 @@
     "./method-override": "./src/middleware/method-override/index.ts",
     "./powered-by": "./src/middleware/powered-by/index.ts",
     "./pretty-json": "./src/middleware/pretty-json/index.ts",
+    "./request-id": "./src/middleware/request-id/request-id.ts",
     "./secure-headers": "./src/middleware/secure-headers/secure-headers.ts",
     "./ssg": "./src/helper/ssg/index.ts",
     "./streaming": "./src/helper/streaming/index.ts",
diff --git a/package.json b/package.json
index 4a1cd916..1a3417bd 100644
--- a/package.json
+++ b/package.json
@@ -223,6 +223,11 @@
       "import": "./dist/middleware/pretty-json/index.js",
       "require": "./dist/cjs/middleware/pretty-json/index.js"
     },
+    "./request-id": {
+      "types": "./dist/types/middleware/request-id/index.d.ts",
+      "import": "./dist/middleware/request-id/index.js",
+      "require": "./dist/cjs/middleware/request-id/index.js"
+    },
     "./secure-headers": {
       "types": "./dist/types/middleware/secure-headers/index.d.ts",
       "import": "./dist/middleware/secure-headers/index.js",
@@ -476,6 +481,9 @@
       "pretty-json": [
         "./dist/types/middleware/pretty-json"
       ],
+      "request-id": [
+        "./dist/types/middleware/request-id"
+      ],
       "streaming": [
         "./dist/types/helper/streaming"
       ],
diff --git a/src/middleware/request-id/index.test.ts b/src/middleware/request-id/index.test.ts
new file mode 100644
index 00000000..739cdded
--- /dev/null
+++ b/src/middleware/request-id/index.test.ts
@@ -0,0 +1,155 @@
+import type { Context } from '../../context'
+import { Hono } from '../../hono'
+import { requestId } from '.'
+
+const regexUUIDv4 = /([0-9a-f]{8})-([0-9a-f]{4})-([0-9a-f]{4})-([0-9a-f]{4})-([0-9a-f]{12})/
+
+describe('Request ID Middleware', () => {
+  const app = new Hono()
+  app.use('*', requestId())
+  app.get('/requestId', (c) => c.text(c.get('requestId') ?? 'No Request ID'))
+
+  it('Should return random request id', async () => {
+    const res = await app.request('http://localhost/requestId')
+    expect(res).not.toBeNull()
+    expect(res.status).toBe(200)
+    expect(res.headers.get('X-Request-Id')).toMatch(regexUUIDv4)
+    expect(await res.text()).match(regexUUIDv4)
+  })
+
+  it('Should return custom request id', async () => {
+    const res = await app.request('http://localhost/requestId', {
+      headers: {
+        'X-Request-Id': 'hono-is-cool',
+      },
+    })
+    expect(res).not.toBeNull()
+    expect(res.status).toBe(200)
+    expect(res.headers.get('X-Request-Id')).toBe('hono-is-cool')
+    expect(await res.text()).toBe('hono-is-cool')
+  })
+
+  it('Should return random request id without using request header', async () => {
+    const res = await app.request('http://localhost/requestId', {
+      headers: {
+        'X-Request-Id': 'Hello!12345-@*^',
+      },
+    })
+    expect(res).not.toBeNull()
+    expect(res.status).toBe(200)
+    expect(res.headers.get('X-Request-Id')).toMatch(regexUUIDv4)
+    expect(await res.text()).toMatch(regexUUIDv4)
+  })
+})
+
+describe('Request ID Middleware with custom generator', () => {
+  function generateWord() {
+    return 'HonoIsWebFramework'
+  }
+  function generateDoubleRequestId(c: Context) {
+    const honoId = c.req.header('Hono-Request-Id')
+    const ohnoId = c.req.header('Ohno-Request-Id')
+    if (honoId && ohnoId) {
+      return honoId + ohnoId
+    }
+    return crypto.randomUUID()
+  }
+  const app = new Hono()
+  app.use('/word', requestId({ generator: generateWord }))
+  app.use('/doubleRequestId', requestId({ generator: generateDoubleRequestId }))
+  app.get('/word', (c) => c.text(c.get('requestId') ?? 'No Request ID'))
+  app.get('/doubleRequestId', (c) => c.text(c.get('requestId') ?? 'No Request ID'))
+  it('Should return custom request id', async () => {
+    const res = await app.request('http://localhost/word')
+    expect(res).not.toBeNull()
+    expect(res.status).toBe(200)
+    expect(res.headers.get('X-Request-Id')).toBe('HonoIsWebFramework')
+    expect(await res.text()).toBe('HonoIsWebFramework')
+  })
+
+  it('Should return complex request id', async () => {
+    const res = await app.request('http://localhost/doubleRequestId', {
+      headers: {
+        'Hono-Request-Id': 'Hello',
+        'Ohno-Request-Id': 'World',
+      },
+    })
+    expect(res).not.toBeNull()
+    expect(res.status).toBe(200)
+    expect(res.headers.get('X-Request-Id')).toBe('HelloWorld')
+    expect(await res.text()).toBe('HelloWorld')
+  })
+})
+
+describe('Request ID Middleware with limit length', () => {
+  const charactersOf255 = 'h'.repeat(255)
+  const charactersOf256 = 'h'.repeat(256)
+
+  const app = new Hono()
+  app.use('/requestId', requestId())
+  app.use('/limit256', requestId({ limitLength: 256 }))
+  app.get('/requestId', (c) => c.text(c.get('requestId') ?? 'No Request ID'))
+  app.get('/limit256', (c) => c.text(c.get('requestId') ?? 'No Request ID'))
+
+  it('Should return custom request id', async () => {
+    const res = await app.request('http://localhost/requestId', {
+      headers: {
+        'X-Request-Id': charactersOf255,
+      },
+    })
+    expect(res).not.toBeNull()
+    expect(res.status).toBe(200)
+    expect(res.headers.get('X-Request-Id')).toBe(charactersOf255)
+    expect(await res.text()).toBe(charactersOf255)
+  })
+  it('Should return random request id without using request header', async () => {
+    const res = await app.request('http://localhost/requestId', {
+      headers: {
+        'X-Request-Id': charactersOf256,
+      },
+    })
+    expect(res).not.toBeNull()
+    expect(res.status).toBe(200)
+    expect(res.headers.get('X-Request-Id')).toMatch(regexUUIDv4)
+    expect(await res.text()).toMatch(regexUUIDv4)
+  })
+  it('Should return custom request id with 256 characters', async () => {
+    const res = await app.request('http://localhost/limit256', {
+      headers: {
+        'X-Request-Id': charactersOf256,
+      },
+    })
+    expect(res).not.toBeNull()
+    expect(res.status).toBe(200)
+    expect(res.headers.get('X-Request-Id')).toBe(charactersOf256)
+    expect(await res.text()).toBe(charactersOf256)
+  })
+})
+
+describe('Request ID Middleware with custom header', () => {
+  const app = new Hono()
+  app.use('/requestId', requestId({ headerName: 'Hono-Request-Id' }))
+  app.get('/emptyId', requestId({ headerName: '' }))
+  app.get('/requestId', (c) => c.text(c.get('requestId') ?? 'No Request ID'))
+  app.get('/emptyId', (c) => c.text(c.get('requestId') ?? 'No Request ID'))
+
+  it('Should return custom request id', async () => {
+    const res = await app.request('http://localhost/requestId', {
+      headers: {
+        'Hono-Request-Id': 'hono-is-cool',
+      },
+    })
+    expect(res).not.toBeNull()
+    expect(res.status).toBe(200)
+    expect(res.headers.get('Hono-Request-Id')).toBe('hono-is-cool')
+    expect(await res.text()).toBe('hono-is-cool')
+  })
+
+  it('Should not return request id', async () => {
+    const res = await app.request('http://localhost/emptyId')
+    expect(res).not.toBeNull()
+    expect(res.status).toBe(200)
+    expect(res.headers.get('X-Request-Id')).toBeNull()
+    expect(await res.text()).toMatch(regexUUIDv4)
+  })
+})
diff --git a/src/middleware/request-id/index.ts b/src/middleware/request-id/index.ts
new file mode 100644
index 00000000..040ec661
--- /dev/null
+++ b/src/middleware/request-id/index.ts
@@ -0,0 +1,8 @@
+import type { RequestIdVariables } from './request-id'
+export type { RequestIdVariables }
+export { requestId } from './request-id'
+import type {} from '../..'
+
+declare module '../..' {
+  interface ContextVariableMap extends RequestIdVariables {}
+}
diff --git a/src/middleware/request-id/request-id.ts b/src/middleware/request-id/request-id.ts
new file mode 100644
index 00000000..f5ae5846
--- /dev/null
+++ b/src/middleware/request-id/request-id.ts
@@ -0,0 +1,59 @@
+/**
+ * @module
+ * Request ID Middleware for Hono.
+ */
+
+import type { Context } from '../../context'
+import type { MiddlewareHandler } from '../../types'
+
+export type RequestIdVariables = {
+  requestId: string
+}
+
+export type RequestIdOptions = {
+  limitLength?: number
+  headerName?: string
+  generator?: (c: Context) => string
+}
+
+/**
+ * Request ID Middleware for Hono.
+ *
+ * @param {object} options - Options for Request ID middleware.
+ * @param {number} [options.limitLength=255] - The maximum length of request id.
+ * @param {string} [options.headerName=X-Request-Id] - The header name used in request id.
+ * @param {generator} [options.generator=() => crypto.randomUUID()] - The request id generation function.
+ *
+ * @returns {MiddlewareHandler} The middleware handler function.
+ *
+ * @example
+ * ```ts
+ * type Variables = RequestIdVariables
+ * const app = new Hono<{Variables: Variables}>()
+ *
+ * app.use(requestId())
+ * app.get('/', (c) => {
+ *   console.log(c.get('requestId')) // Debug
+ *   return c.text('Hello World!')
+ * })
+ * ```
+ */
+export const requestId = ({
+  limitLength = 255,
+  headerName = 'X-Request-Id',
+  generator = () => crypto.randomUUID(),
+}: RequestIdOptions = {}): MiddlewareHandler => {
+  return async function requestId(c, next) {
+    // If `headerName` is empty string, req.header will return the object
+    let reqId = headerName ? c.req.header(headerName) : undefined
+    if (!reqId || reqId.length > limitLength || /[^\w\-]/.test(reqId)) {
+      reqId = generator(c)
+    }
+
+    c.set('requestId', reqId)
+    if (headerName) {
+      c.header(headerName, reqId)
+    }
+    await next()
+  }
+}