From fdf77862acbbfc1540ff0732a4cfd74a55781393 Mon Sep 17 00:00:00 2001
From: Marcel Overdijk <marceloverdijk@gmail.com>
Date: Sun, 8 Sep 2024 08:53:01 +0200
Subject: [PATCH] feat(contextStorage): added Context Storage Middleware
 (#3373)

---
 jsr.json                                     |  1 +
 package.json                                 |  8 +++
 src/middleware/context-storage/index.test.ts | 30 +++++++++++
 src/middleware/context-storage/index.ts      | 55 ++++++++++++++++++++
 4 files changed, 94 insertions(+)
 create mode 100644 src/middleware/context-storage/index.test.ts
 create mode 100644 src/middleware/context-storage/index.ts

diff --git a/jsr.json b/jsr.json
index bfeae8b2..e6af4ed6 100644
--- a/jsr.json
+++ b/jsr.json
@@ -26,6 +26,7 @@
     "./cookie": "./src/helper/cookie/index.ts",
     "./accepts": "./src/helper/accepts/index.ts",
     "./compress": "./src/middleware/compress/index.ts",
+    "./context-storage": "./src/middleware/context-storage/index.ts",
     "./cors": "./src/middleware/cors/index.ts",
     "./csrf": "./src/middleware/csrf/index.ts",
     "./etag": "./src/middleware/etag/index.ts",
diff --git a/package.json b/package.json
index f7647825..363aaaff 100644
--- a/package.json
+++ b/package.json
@@ -104,6 +104,11 @@
       "import": "./dist/middleware/compress/index.js",
       "require": "./dist/cjs/middleware/compress/index.js"
     },
+    "./context-storage": {
+      "types": "./dist/types/middleware/context-storage/index.d.ts",
+      "import": "./dist/middleware/context-storage/index.js",
+      "require": "./dist/cjs/middleware/context-storage/index.js"
+    },
     "./cors": {
       "types": "./dist/types/middleware/cors/index.d.ts",
       "import": "./dist/middleware/cors/index.js",
@@ -421,6 +426,9 @@
       "compress": [
         "./dist/types/middleware/compress"
       ],
+      "context-storage": [
+        "./dist/types/middleware/context-storage"
+      ],
       "cors": [
         "./dist/types/middleware/cors"
       ],
diff --git a/src/middleware/context-storage/index.test.ts b/src/middleware/context-storage/index.test.ts
new file mode 100644
index 00000000..d9713b97
--- /dev/null
+++ b/src/middleware/context-storage/index.test.ts
@@ -0,0 +1,30 @@
+import { Hono } from '../../hono'
+import { contextStorage, getContext } from '.'
+
+describe('Context Storage Middleware', () => {
+  type Env = {
+    Variables: {
+      message: string
+    }
+  }
+
+  const app = new Hono<Env>()
+
+  app.use(contextStorage())
+  app.use(async (c, next) => {
+    c.set('message', 'Hono is cool!!')
+    await next()
+  })
+  app.get('/', (c) => {
+    return c.text(getMessage())
+  })
+
+  const getMessage = () => {
+    return getContext<Env>().var.message
+  }
+
+  it('Should get context', async () => {
+    const res = await app.request('/')
+    expect(await res.text()).toBe('Hono is cool!!')
+  })
+})
diff --git a/src/middleware/context-storage/index.ts b/src/middleware/context-storage/index.ts
new file mode 100644
index 00000000..2545d3ca
--- /dev/null
+++ b/src/middleware/context-storage/index.ts
@@ -0,0 +1,55 @@
+/**
+ * @module
+ * Context Storage Middleware for Hono.
+ */
+
+import type { Context } from '../../context'
+import type { Env, MiddlewareHandler } from '../../types'
+import { AsyncLocalStorage } from 'node:async_hooks'
+
+const asyncLocalStorage = new AsyncLocalStorage<Context>()
+
+/**
+ * Context Storage Middleware for Hono.
+ *
+ * @see {@link https://hono.dev/docs/middleware/builtin/context-storage}
+ *
+ * @returns {MiddlewareHandler} The middleware handler function.
+ *
+ * @example
+ * ```ts
+ * type Env = {
+ *   Variables: {
+ *     message: string
+ *   }
+ * }
+ *
+ * const app = new Hono<Env>()
+ *
+ * app.use(contextStorage())
+ *
+ * app.use(async (c, next) => {
+ *   c.set('message', 'Hono is cool!!)
+ *   await next()
+ * })
+ *
+ * app.get('/', async (c) => { c.text(getMessage()) })
+ *
+ * const getMessage = () => {
+ *   return getContext<Env>().var.message
+ * }
+ * ```
+ */
+export const contextStorage = (): MiddlewareHandler => {
+  return async function contextStorage(c, next) {
+    await asyncLocalStorage.run(c, next)
+  }
+}
+
+export const getContext = <E extends Env = Env>(): Context<E> => {
+  const context = asyncLocalStorage.getStore()
+  if (!context) {
+    throw new Error('Context is not available')
+  }
+  return context
+}