feat(middleware): Add .shortCircuit to MiddlewareResponse (#10586)

### `MiddlewaresResponse.shortCircuit()`

Adds a helper to generate a intercept/short-circuit response, that will
interrupt execution of _all_ middleware and react rendering, and
immediately return the response.

There's a few different ways you can use this, see examples below:

```ts
const shortCircuitMw: Middleware = (req, res) => {
  // A) You can short circuit after building the response (or use the res param)
  // This allows you to use all the convenience helpers like cookies of MW Response
  if (req.url.includes('create-new-response')) {
    const shortCircuitResponse = new MiddlewareResponse('Short-circuiting')
    shortCircuitResponse.headers.set('shortCircuit', 'yes')
    shortCircuitResponse.cookies.set('shortCircuitCookie', 'do-not-allow', {
      expires: new Date(Date.now() + 1000 * 60 * 60),
    })
    shortCircuitResponse.shortCircuit()
  }

  // B) You can directly construct a new short-circuit response
  // (discarding whatever response was built before)
  if (req.url.includes('using-existing-res')) {
    res.shortCircuit('Short-circuiting directly', {
      headers: { shortCircuitDirect: 'yes' },
    })
  }

}
```

.changesets/10586.md

@@ -0,0 +1,30 @@
+- feat(middleware): Add .shortCircuit to MiddlewareResponse (#10586) by @dac09
+
+Adds a helper to generate a intercept/short-circuit response, that will interrupt execution of _all_ middleware and react rendering, and immediately return the response. 
+
+There's a few different ways you can use this, see examples below:
+
+```ts
+const shortCircuitMw: Middleware = (req, res) => {
+  // A) You can short circuit after building the response (or use the res param)
+  // This allows you to use all the convenience helpers like cookies of MW Response
+  if (req.url.includes('create-new-response')) {
+    const shortCircuitResponse = new MiddlewareResponse('Short-circuiting')
+    shortCircuitResponse.headers.set('shortCircuit', 'yes')
+    shortCircuitResponse.cookies.set('shortCircuitCookie', 'do-not-allow', {
+      expires: new Date(Date.now() + 1000 * 60 * 60),
+    })
+    shortCircuitResponse.shortCircuit()
+  }
+
+  // B) You can directly construct a new short-circuit response
+  // (discarding whatever response was built before)
+  if (req.url.includes('using-existing-res')) {
+    res.shortCircuit('Short-circuiting directly', {
+      headers: { shortCircuitDirect: 'yes' },
+    })
+  }
+
+}
+```
+

packages/vite/src/middleware/MiddlewareResponse.test.ts

@@ -1,7 +1,10 @@
 import { Response as PonyfillResponse } from '@whatwg-node/fetch'
 import { describe, expect, test } from 'vitest'
 
-import { MiddlewareResponse } from './MiddlewareResponse'
+import {
+  MiddlewareResponse,
+  MiddlewareShortCircuit,
+} from './MiddlewareResponse'
 
 describe('MiddlewareResponse', () => {
   test('constructor', () => {
@@ -73,4 +76,53 @@ describe('MiddlewareResponse', () => {
       '/bye',
     )
   })
+
+  test('Constructs short-circuits correctly when parameters passed to it', async () => {
+    const myMwResponse = new MiddlewareResponse()
+
+    try {
+      myMwResponse.shortCircuit(JSON.stringify({ shortCircuit: true }), {
+        status: 401,
+      })
+    } catch (e) {
+      const shortCircuitError = e as MiddlewareShortCircuit
+      expect(shortCircuitError instanceof MiddlewareShortCircuit).toBe(true)
+      expect(shortCircuitError.mwResponse.toResponse().status).toStrictEqual(
+        401,
+      )
+      expect(
+        await shortCircuitError.mwResponse.toResponse().json(),
+      ).toStrictEqual({
+        shortCircuit: true,
+      })
+    }
+
+    expect.assertions(3)
+  })
+
+  test('Constructs short-circuits using existing response properties when parameters passed to it', async () => {
+    const myMwResponse = new MiddlewareResponse('Nope', {
+      status: 429,
+      statusText: 'Hold your horses!',
+    })
+
+    try {
+      myMwResponse.shortCircuit()
+    } catch (e) {
+      const shortCircuitError = e as MiddlewareShortCircuit
+      expect(shortCircuitError instanceof MiddlewareShortCircuit).toBe(true)
+      expect(shortCircuitError.mwResponse.toResponse().status).toStrictEqual(
+        429,
+      )
+      expect(
+        await shortCircuitError.mwResponse.toResponse().text(),
+      ).toStrictEqual('Nope')
+
+      expect(
+        await shortCircuitError.mwResponse.toResponse().statusText,
+      ).toStrictEqual('Hold your horses!')
+    }
+
+    expect.assertions(4)
+  })
 })

packages/vite/src/middleware/MiddlewareResponse.ts

@@ -3,6 +3,16 @@ import cookie from 'cookie'
 
 import { CookieJar } from './CookieJar.js'
 
+export class MiddlewareShortCircuit extends Error {
+  mwResponse: MiddlewareResponse
+
+  constructor(body?: BodyInit | null, responseInit?: ResponseInit) {
+    super('Short cirtcuit. Skipping all middleware, and returning early')
+    this.name = 'MiddlewareShortCircuit'
+    this.mwResponse = new MiddlewareResponse(body, responseInit)
+  }
+}
+
 /**
  * This is actually a Response builder class
  * After setting all the required proeprties, we can call `build` to get a Web API Response object
@@ -12,24 +22,63 @@ export class MiddlewareResponse {
   headers = new Headers()
   body: BodyInit | null | undefined
   status = 200
+  statusText: string | undefined
 
   constructor(body?: BodyInit | null, init?: ResponseInit) {
     this.body = body
     this.headers = new Headers(init?.headers)
     this.status = init?.status || 200
+    this.statusText = init?.statusText
   }
 
   static fromResponse = (res: Response) => {
     return new MiddlewareResponse(res.body, {
       headers: res.headers,
       status: res.status,
+      statusText: res.statusText,
     })
   }
 
+  /**
+   *
+   * Short circuit the middleware chain and return early.
+   * This will skip all the remaining middleware and return the response immediately.
+   *
+   * @returns MiddlewareResponse
+   */
+  shortCircuit = (body?: BodyInit | null, init?: ResponseInit) => {
+    for (const [ckName, ckParams] of this.cookies) {
+      this.headers.append(
+        'Set-Cookie',
+        cookie.serialize(ckName, ckParams.value, ckParams.options),
+      )
+    }
+
+    throw new MiddlewareShortCircuit(
+      body || this.body,
+      init || {
+        headers: this.headers,
+        status: this.status,
+        statusText: this.statusText,
+      },
+    )
+  }
+
+  /**
+   * Skip the current middleware and move to the next one.
+   * Careful: It creates a new Response, so any middleware that modifies the response before the current one will be lost.
+   * @returns MiddlewareResponse
+   */
   static next = () => {
     return new MiddlewareResponse()
   }
 
+  /**
+   *
+   * Return a MiddlewareResponse object that will redirect the client to the specified location
+   *
+   * @returns MiddlewareResponse
+   */
   static redirect = (
     location: string,
     type: 'permanent' | 'temporary' = 'temporary',
@@ -56,6 +105,7 @@ export class MiddlewareResponse {
     return new PonyResponse(this.body, {
       headers: this.headers,
       status: this.status,
+      statusText: this.statusText,
     })
   }
 }

packages/vite/src/middleware/invokeMiddleware.test.ts

@@ -8,6 +8,7 @@ import { createServerStorage } from '../serverStore'
 import { invoke } from './invokeMiddleware'
 import type { MiddlewareRequest } from './MiddlewareRequest'
 import { MiddlewareResponse } from './MiddlewareResponse'
+import type { Middleware } from './types'
 
 describe('Invoke middleware', () => {
   beforeAll(() => {
@@ -50,6 +51,7 @@ describe('Invoke middleware', () => {
       consoleErrorSpy.mockRestore()
     })
 
+    // This means that will CONTINUE execution of the middleware chain, and react rendering
     test('returns a MiddlewareResponse, even if middleware throws', async () => {
       const throwingMiddleware = () => {
         throw new Error('I want to break free')
@@ -63,5 +65,57 @@ describe('Invoke middleware', () => {
       expect(mwRes).toBeInstanceOf(MiddlewareResponse)
       expect(authState).toEqual(middlewareDefaultAuthProviderState)
     })
+
+    // A short-circuit is a way to stop the middleware chain immediately, and return a response
+    test('will return a MiddlewareResposne, if a short-circuit is thrown', async () => {
+      const shortCircuitMW: Middleware = (_req, res) => {
+        res.shortCircuit('Zap', {
+          status: 999,
+          statusText: 'Ouch',
+        })
+      }
+
+      const [mwRes, authState] = await invoke(
+        new Request('https://example.com'),
+        shortCircuitMW,
+      )
+
+      expect(mwRes).toBeInstanceOf(MiddlewareResponse)
+      expect(mwRes.body).toEqual('Zap')
+      expect(mwRes.status).toEqual(999)
+      expect(mwRes.statusText).toEqual('Ouch')
+      expect(authState).toEqual(middlewareDefaultAuthProviderState)
+    })
+
+    test('can set extra properties in the shortcircuit response', async () => {
+      const testMw: Middleware = () => {
+        const shortCircuitRes = new MiddlewareResponse('Zap')
+
+        shortCircuitRes.cookies.set('monster', 'nomnomnom', {
+          expires: new Date('2022-01-01'),
+        })
+        shortCircuitRes.headers.set('redwood', 'is awesome')
+
+        shortCircuitRes.shortCircuit()
+      }
+
+      const [mwRes, authState] = await invoke(
+        new Request('https://example.com'),
+        testMw,
+      )
+
+      expect(mwRes).toBeInstanceOf(MiddlewareResponse)
+      expect(mwRes.body).toEqual('Zap')
+      expect(mwRes.status).toEqual(200)
+
+      expect(mwRes.toResponse().headers.getSetCookie()).toContainEqual(
+        'monster=nomnomnom; Expires=Sat, 01 Jan 2022 00:00:00 GMT',
+      )
+      expect(mwRes.toResponse().headers.get('redwood')).toStrictEqual(
+        'is awesome',
+      )
+
+      expect(authState).toEqual(middlewareDefaultAuthProviderState)
+    })
   })
 })

packages/vite/src/middleware/invokeMiddleware.ts

@@ -6,7 +6,10 @@ import {
 import { setServerAuthState } from '../serverStore.js'
 
 import { MiddlewareRequest } from './MiddlewareRequest.js'
-import { MiddlewareResponse } from './MiddlewareResponse.js'
+import {
+  MiddlewareResponse,
+  MiddlewareShortCircuit,
+} from './MiddlewareResponse.js'
 import type { Middleware, MiddlewareInvokeOptions } from './types.js'
 
 /**
@@ -49,6 +52,12 @@ export const invoke = async (
       )
     }
   } catch (e) {
+    // @TODO catch the error here, and see if its a short-circuit
+    // A shortcircuit will prevent execution of all other middleware down the chain, and prevent react rendering
+    if (e instanceof MiddlewareShortCircuit) {
+      return [e.mwResponse, mwReq.serverAuthContext.get()]
+    }
+
     console.error('Error executing middleware > \n')
     console.error('~'.repeat(80))
     console.error(e)