fix: Return MiddlewareResponse obj for rewrite (#2159)

* fix: Return MiddlewareResponse for rewrite

* chore: remove console logs

* chore: added tests and updated logic to include init

* chore: revert selector back to original

* fix: pass in new headers

* chore: applyheader before response

* chore: update tests remove header fix

* chore: testing update

* chore: removed undici import

* chore: added undici import back

* Update demos/middleware/middleware.ts

Co-authored-by: Nick Taylor <[email protected]>

* Update cypress/e2e/middleware/enhanced.cy.ts

update pageProps testing

Co-authored-by: Nick Taylor <[email protected]>

* chore: removed header override testing

* chore: formatting

* chore: new line

---------

Co-authored-by: Nick Taylor <[email protected]>

Author: taty2010

cypress/e2e/middleware/enhanced.cy.ts

@@ -1,5 +1,19 @@
 describe('Enhanced middleware', () => {
-  it('rewrites the response body', () => {
+  it('rewrites the response body using request.rewrite()', () => {
+    cy.visit('/request-rewrite')
+    cy.get('#message').contains('This was static (& escaping test &) but has been transformed in')
+    cy.contains("This is an ad that isn't shown by default")
+  })
+
+  it('modifies the page props when using request.rewrite()', () => {
+    cy.visit('/request-rewrite')
+    cy.get('script#__NEXT_DATA__').then((element) => {
+      const { props } = JSON.parse(element.text());
+      expect(props.pageProps.message).to.include('This was static (& escaping test &) but has been transformed in')
+    })
+  })
+
+  it('rewrites the response body using request.next()', () => {
     cy.visit('/static')
     cy.get('#message').contains('This was static (& escaping test &) but has been transformed in')
     cy.contains("This is an ad that isn't shown by default")

demos/middleware/middleware.ts

@@ -27,7 +27,22 @@ export async function middleware(req: NextRequest) {
     const message = `This was static (& escaping test &) but has been transformed in ${req.geo?.city}`
 
     // Transform the response HTML and props
-    res.replaceText('p[id=message]', message)
+    res.replaceText('#message', message)
+    res.setPageProp('message', message)
+    res.setPageProp('showAd', true)
+
+    res.headers.set('x-modified-edge', 'true')
+    res.headers.set('x-is-deno', 'Deno' in globalThis ? 'true' : 'false')
+    return res
+  }
+
+  if (pathname.startsWith('/request-rewrite')) {
+    // request.rewrite() should return the MiddlewareResponse object instead of the Response object.
+    const res = await request.rewrite('/static-rewrite')
+    const message = `This was static (& escaping test &) but has been transformed in ${req.geo?.city}`
+
+    // Transform the response HTML and props
+    res.replaceText('#message', message)
     res.setPageProp('message', message)
     res.setPageProp('showAd', true)
 
@@ -138,6 +153,7 @@ export const config = {
     '/headers',
     '/cookies/:path*',
     { source: '/static' },
+    {source: '/request-rewrite' },
     { source: '/matcher-cookie'},
     { source: '/shows/((?!99|88).*)' },
     {

demos/middleware/pages/request-rewrite.js

@@ -0,0 +1,10 @@
+const Rewrite = () => {
+    return (
+      <div>
+        <p>This should have been rewritten</p>
+      </div>
+    )
+  }
+  
+export default Rewrite
+

demos/middleware/pages/static-rewrite.js

@@ -0,0 +1,37 @@
+import * as React from 'react'
+
+const useHydrated = () => {
+  const [hydrated, setHydrated] = React.useState(false)
+  React.useEffect(() => {
+    setHydrated(true)
+  }, [])
+  return hydrated
+}
+
+const Page = ({ message, showAd }) => {
+  const hydrated = useHydrated()
+  return (
+    <div>
+      <p id="message">{message}</p>
+      {hydrated && showAd ? (
+        <div>
+          <p>This is an ad that isn't shown by default on static test 2 page</p>
+          <img src="http://placekitten.com/400/300" />
+        </div>
+      ) : (
+        <p>No ads for me</p>
+      )}
+    </div>
+  )
+}
+
+export async function getStaticProps() {
+  return {
+    props: {
+      message: 'This is a static page',
+      showAd: false,
+    },
+  }
+}
+
+export default Page

packages/next/src/middleware/request.ts

@@ -1,6 +1,5 @@
 import type { Context } from '@netlify/edge-functions'
 import type { NextURL } from 'next/dist/server/web/next-url'
-import { NextResponse } from 'next/server'
 import type { NextRequest as InternalNextRequest } from 'next/server'
 
 import { MiddlewareResponse } from './response'
@@ -64,16 +63,17 @@ export class MiddlewareRequest extends Request {
     if (response.status === 301 && locationHeader?.startsWith('/')) {
       response = await this.context.rewrite(locationHeader)
     }
-
     return new MiddlewareResponse(response)
   }
 
-  rewrite(destination: string | URL | NextURL, init?: ResponseInit): NextResponse {
+  async rewrite(destination: string | URL | NextURL, init?: ResponseInit) {
     if (typeof destination === 'string' && destination.startsWith('/')) {
       destination = new URL(destination, this.url)
     }
     this.applyHeaders()
-    return NextResponse.rewrite(destination, init)
+    const response = await this.context.rewrite(destination)
+
+    return new MiddlewareResponse(response, init)
   }
 
   get headers() {

packages/next/src/middleware/response.ts

@@ -10,7 +10,9 @@ export type NextDataTransform = <T extends { pageProps?: Record<string, any> }>(
 export class MiddlewareResponse extends NextResponse {
   private readonly dataTransforms: NextDataTransform[]
   private readonly elementHandlers: Array<[selector: string, handlers: ElementHandlers]>
-  constructor(public originResponse: Response) {
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  constructor(public originResponse: Response, init?: ResponseInit) {
     // we need to propagate the set-cookie header, so response.cookies.get works correctly
     const initHeaders = new Headers()
     if (originResponse.headers.has('set-cookie')) {

test/e2e/next-test-lib/next-test-utils.js

@@ -2,7 +2,7 @@
 import spawn from 'cross-spawn'
 import { existsSync, readFileSync, unlinkSync, writeFileSync } from 'fs'
 import { writeFile } from 'fs-extra'
-import { fetch as undiciFetch } from 'undici'
+import { fetch as undiciFetch } from 'next/dist/compiled/undici'
 import nodeFetch from 'node-fetch'
 import path from 'path'
 import qs from 'querystring'
@@ -111,7 +111,7 @@ async function processChunkedResponse(response) {
  * @param {string | number} appPort
  * @param {string} pathname
  * @param {Record<string, any> | string | undefined} [query]
- * @param {import("undici").RequestInit} [opts]
+ * @param {import("next/dist/compiled/undici").RequestInit} [opts]
  * @returns {Promise<string>}
  */
 export function renderViaHTTP(appPort, pathname, query, opts) {
@@ -121,11 +121,11 @@ export function renderViaHTTP(appPort, pathname, query, opts) {
 /**
  * @param {string | number} appPort
  * @param {string} pathname
- * @param {Record<string, any> | string | undefined} query
- * @param {RequestInit} opts
- * @returns {Promise<Response>}
+ * @param {Record<string, any> | string | null | undefined} [query]
+ * @param {import('node-fetch').RequestInit} [opts]
+ * @returns {Promise<Response & {buffer: any} & {headers: Headers}>}
  */
-export async function fetchViaHTTP(appPort, pathname, query = undefined, opts = undefined, useUndici = false) {
+export function fetchViaHTTP(appPort, pathname, query, opts, useUndici = false) {
   const url = `${pathname}${typeof query === 'string' ? query : query ? `?${qs.stringify(query)}` : ''}`
   const fetch = useUndici ? undiciFetch : nodeFetch
   const fullUrl = getFullUrl(appPort, url)