Implement MutableRequestCookies in server entries (#48847)

Similar to #47922 but based off the latest server implementation and #48626:

> This PR implements the MutableRequestCookies instance for cookies() based on the current async context, so we can allow setting cookies in certain places such as Server Functions and Route handlers. Note that to support Route Handlers, we need to also implement the logic of merging Response's Set-Cookie header and the cookies() mutations, hence it's not included in this PR.
>
> fix [NEXT-942](https://linear.app/vercel/issue/NEXT-942)

This PR also adds the same support for Custom Routes.

cc @styfle.

fix NEXT-942, fix NEXT-941.
This commit is contained in:
Shu Ding 2023-04-26 15:19:01 +02:00 committed by GitHub
parent da2804f974
commit b21fd96606
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 318 additions and 121 deletions

View file

@ -1,47 +1,11 @@
import type { AsyncLocalStorage } from 'async_hooks'
import { createAsyncLocalStorage } from './async-local-storage'
export interface ActionStore {
readonly isAction: boolean
readonly isAction?: boolean
readonly isAppRoute?: boolean
}
export type ActionAsyncStorage = AsyncLocalStorage<ActionStore>
let createAsyncLocalStorage: () => ActionAsyncStorage
if (process.env.NEXT_RUNTIME === 'edge') {
createAsyncLocalStorage = () => {
let store: ActionStore | undefined
return {
disable() {
throw new Error(
'Invariant: AsyncLocalStorage accessed in runtime where it is not available'
)
},
getStore() {
return store
},
async run<R>(s: ActionStore, fn: () => R): Promise<R> {
store = s
try {
return await fn()
} finally {
store = undefined
}
},
exit() {
throw new Error(
'Invariant: AsyncLocalStorage accessed in runtime where it is not available'
)
},
enterWith() {
throw new Error(
'Invariant: AsyncLocalStorage accessed in runtime where it is not available'
)
},
} as ActionAsyncStorage
}
} else {
createAsyncLocalStorage =
require('./async-local-storage').createAsyncLocalStorage
}
export const actionAsyncStorage: ActionAsyncStorage = createAsyncLocalStorage()

View file

@ -2,6 +2,7 @@ import { RequestCookiesAdapter } from '../../server/web/spec-extension/adapters/
import { HeadersAdapter } from '../../server/web/spec-extension/adapters/headers'
import { RequestCookies } from '../../server/web/spec-extension/cookies'
import { requestAsyncStorage } from './request-async-storage'
import { actionAsyncStorage } from './action-async-storage'
import { staticGenerationBailout } from './static-generation-bailout'
export function headers() {
@ -42,5 +43,13 @@ export function cookies() {
)
}
const asyncActionStore = actionAsyncStorage.getStore()
if (
asyncActionStore &&
(asyncActionStore.isAction || asyncActionStore.isAppRoute)
) {
return requestStore.mutableCookies
}
return requestStore.cookies
}

View file

@ -1,4 +1,5 @@
import type { AsyncLocalStorage } from 'async_hooks'
import type { RequestCookies } from '../../server/web/spec-extension/cookies'
import type { PreviewData } from '../../../types'
import type { ReadonlyHeaders } from '../../server/web/spec-extension/adapters/headers'
import type { ReadonlyRequestCookies } from '../../server/web/spec-extension/adapters/request-cookies'
@ -8,6 +9,7 @@ import { createAsyncLocalStorage } from './async-local-storage'
export interface RequestStore {
readonly headers: ReadonlyHeaders
readonly cookies: ReadonlyRequestCookies
readonly mutableCookies: RequestCookies
readonly previewData: PreviewData
}

View file

@ -14,6 +14,7 @@ import {
type ReadonlyHeaders,
} from '../web/spec-extension/adapters/headers'
import {
MutableRequestCookiesAdapter,
RequestCookiesAdapter,
type ReadonlyRequestCookies,
} from '../web/spec-extension/adapters/request-cookies'
@ -35,6 +36,14 @@ function getCookies(
return RequestCookiesAdapter.seal(cookies)
}
function getMutableCookies(
headers: Headers | IncomingHttpHeaders,
res: ServerResponse | BaseNextResponse | undefined
): RequestCookies {
const cookies = new RequestCookies(HeadersAdapter.from(headers))
return MutableRequestCookiesAdapter.seal(cookies, res)
}
/**
* Tries to get the preview data on the request for the given route. This
* isn't enabled in the edge runtime yet.
@ -80,6 +89,7 @@ export const RequestAsyncStorageWrapper: AsyncStorageWrapper<
const cache: {
headers?: ReadonlyHeaders
cookies?: ReadonlyRequestCookies
mutableCookies?: RequestCookies
} = {}
const store: RequestStore = {
@ -101,6 +111,12 @@ export const RequestAsyncStorageWrapper: AsyncStorageWrapper<
return cache.cookies
},
get mutableCookies() {
if (!cache.mutableCookies) {
cache.mutableCookies = getMutableCookies(req.headers, res)
}
return cache.mutableCookies
},
previewData,
}

View file

@ -31,6 +31,9 @@ import { RouteKind } from '../../route-kind'
import * as Log from '../../../../build/output/log'
import { autoImplementMethods } from './helpers/auto-implement-methods'
import { getNonStaticMethods } from './helpers/get-non-static-methods'
import { SYMBOL_MODIFY_COOKIE_VALUES } from '../../../web/spec-extension/adapters/request-cookies'
import { ResponseCookies } from '../../../web/spec-extension/cookies'
import { HeadersAdapter } from '../../../web/spec-extension/adapters/headers'
/**
* AppRouteRouteHandlerContext is the context that is passed to the route
@ -249,93 +252,138 @@ export class AppRouteRouteModule extends RouteModule<
// Run the handler with the request AsyncLocalStorage to inject the helper
// support. We set this to `unknown` because the type is not known until
// runtime when we do a instanceof check below.
const response: unknown = await RequestAsyncStorageWrapper.wrap(
this.requestAsyncStorage,
requestContext,
const response: unknown = await this.actionAsyncStorage.run(
{
isAppRoute: true,
},
() =>
StaticGenerationAsyncStorageWrapper.wrap(
this.staticGenerationAsyncStorage,
staticGenerationContext,
(staticGenerationStore) => {
// Check to see if we should bail out of static generation based on
// having non-static methods.
if (this.nonStaticMethods) {
this.staticGenerationBailout(
`non-static methods used ${this.nonStaticMethods.join(', ')}`
)
}
RequestAsyncStorageWrapper.wrap(
this.requestAsyncStorage,
requestContext,
() =>
StaticGenerationAsyncStorageWrapper.wrap(
this.staticGenerationAsyncStorage,
staticGenerationContext,
(staticGenerationStore) => {
// Check to see if we should bail out of static generation based on
// having non-static methods.
if (this.nonStaticMethods) {
this.staticGenerationBailout(
`non-static methods used ${this.nonStaticMethods.join(
', '
)}`
)
}
// Update the static generation store based on the dynamic property.
switch (this.dynamic) {
case 'force-dynamic':
// The dynamic property is set to force-dynamic, so we should
// force the page to be dynamic.
staticGenerationStore.forceDynamic = true
this.staticGenerationBailout(`force-dynamic`, {
dynamic: this.dynamic,
})
break
case 'force-static':
// The dynamic property is set to force-static, so we should
// force the page to be static.
staticGenerationStore.forceStatic = true
break
case 'error':
// The dynamic property is set to error, so we should throw an
// error if the page is being statically generated.
staticGenerationStore.dynamicShouldError = true
break
default:
break
}
// Update the static generation store based on the dynamic property.
switch (this.dynamic) {
case 'force-dynamic':
// The dynamic property is set to force-dynamic, so we should
// force the page to be dynamic.
staticGenerationStore.forceDynamic = true
this.staticGenerationBailout(`force-dynamic`, {
dynamic: this.dynamic,
})
break
case 'force-static':
// The dynamic property is set to force-static, so we should
// force the page to be static.
staticGenerationStore.forceStatic = true
break
case 'error':
// The dynamic property is set to error, so we should throw an
// error if the page is being statically generated.
staticGenerationStore.dynamicShouldError = true
break
default:
break
}
// If the static generation store does not have a revalidate value
// set, then we should set it the revalidate value from the userland
// module or default to false.
staticGenerationStore.revalidate ??=
this.userland.revalidate ?? false
// If the static generation store does not have a revalidate value
// set, then we should set it the revalidate value from the userland
// module or default to false.
staticGenerationStore.revalidate ??=
this.userland.revalidate ?? false
// Wrap the request so we can add additional functionality to cases
// that might change it's output or affect the rendering.
const wrappedRequest = proxyRequest(
request,
{ dynamic: this.dynamic },
{
headerHooks: this.headerHooks,
serverHooks: this.serverHooks,
staticGenerationBailout: this.staticGenerationBailout,
}
)
// TODO: propagate this pathname from route matcher
const route = getPathnameFromAbsolutePath(this.resolvedPagePath)
getTracer().getRootSpanAttributes()?.set('next.route', route)
return getTracer().trace(
AppRouteRouteHandlersSpan.runHandler,
{
spanName: `executing api route (app) ${route}`,
attributes: {
'next.route': route,
},
},
async () => {
// Patch the global fetch.
patchFetch({
serverHooks: this.serverHooks,
staticGenerationAsyncStorage:
this.staticGenerationAsyncStorage,
})
const res = await handler(wrappedRequest, {
params: context.params,
})
await Promise.all(
staticGenerationStore.pendingRevalidates || []
// Wrap the request so we can add additional functionality to cases
// that might change it's output or affect the rendering.
const wrappedRequest = proxyRequest(
request,
{ dynamic: this.dynamic },
{
headerHooks: this.headerHooks,
serverHooks: this.serverHooks,
staticGenerationBailout: this.staticGenerationBailout,
}
)
// TODO: propagate this pathname from route matcher
const route = getPathnameFromAbsolutePath(this.resolvedPagePath)
getTracer().getRootSpanAttributes()?.set('next.route', route)
return getTracer().trace(
AppRouteRouteHandlersSpan.runHandler,
{
spanName: `executing api route (app) ${route}`,
attributes: {
'next.route': route,
},
},
async () => {
// Patch the global fetch.
patchFetch({
serverHooks: this.serverHooks,
staticGenerationAsyncStorage:
this.staticGenerationAsyncStorage,
})
const res = await handler(wrappedRequest, {
params: context.params,
})
await Promise.all(
staticGenerationStore.pendingRevalidates || []
)
// It's possible cookies were set in the handler, so we need
// to merge the modified cookies and the returned response
// here.
// TODO: Move this into a helper function.
const requestStore = this.requestAsyncStorage.getStore()
if (requestStore && requestStore.mutableCookies) {
const modifiedCookieValues = (
requestStore.mutableCookies as any
)[SYMBOL_MODIFY_COOKIE_VALUES] as [string, string][]
if (modifiedCookieValues.length) {
// Return a new response that extends the response with
// the modified cookies as fallbacks. `res`' cookies
// will still take precedence.
const resCookies = new ResponseCookies(
HeadersAdapter.from(res.headers)
)
const finalCookies = resCookies.getAll()
// Set the modified cookies as fallbacks.
modifiedCookieValues.forEach((cookie) =>
resCookies.set(cookie[0], cookie[1])
)
// Set the original cookies as the final values.
finalCookies.forEach((cookie) => resCookies.set(cookie))
return new Response(res.body, {
status: res.status,
statusText: res.statusText,
headers: {
...res.headers,
'Set-Cookie': resCookies.toString(),
},
})
}
}
return res
}
)
return res
}
)
}
)
)

View file

@ -15,6 +15,8 @@ const headerHooks =
require('next/dist/client/components/headers') as typeof import('../../../client/components/headers')
const { staticGenerationBailout } =
require('next/dist/client/components/static-generation-bailout') as typeof import('../../../client/components/static-generation-bailout')
const { actionAsyncStorage } =
require('next/dist/client/components/action-async-storage') as typeof import('../../../client/components/action-async-storage')
/**
* RouteModuleOptions is the options that are passed to the route module, other
@ -72,6 +74,12 @@ export abstract class RouteModule<
*/
public readonly staticGenerationBailout = staticGenerationBailout
/**
* A reference to the mutation related async storage, such as mutations of
* cookies.
*/
public readonly actionAsyncStorage = actionAsyncStorage
/**
* The userland module. This is the module that is exported from the user's
* code. This is marked as readonly to ensure that the module is not mutated

View file

@ -1,4 +1,7 @@
import type { RequestCookies } from '../cookies'
import type { BaseNextResponse } from '../../../base-http'
import type { ServerResponse } from 'http'
import { ReflectAdapter } from './reflect'
/**
@ -37,3 +40,92 @@ export class RequestCookiesAdapter {
})
}
}
export const SYMBOL_MODIFY_COOKIE_VALUES = Symbol.for('next.mutated.cookies')
export class MutableRequestCookiesAdapter {
public static seal(
cookies: RequestCookies,
res: ServerResponse | BaseNextResponse | undefined
): RequestCookies {
let modifiedValues: [string, string][] = []
const modifiedCookies = new Set<string>()
const updateResponseCookies = () => {
const allCookies = cookies.getAll()
modifiedValues = allCookies
.filter((c) => modifiedCookies.has(c.name))
.map((c) => [c.name, c.value])
if (res) {
res.setHeader(
'Set-Cookie',
modifiedValues.map((c) => `${c[0]}=${c[1]}`)
)
}
}
return new Proxy(cookies, {
get(target, prop, receiver) {
switch (prop) {
// A special symbol to get the modified cookie values
case SYMBOL_MODIFY_COOKIE_VALUES:
return modifiedValues
// TODO: Throw error if trying to set a cookie after the response
// headers have been set.
case 'clear':
return function () {
for (const c of cookies.getAll()) {
modifiedCookies.add(c.name)
}
try {
return cookies.clear()
} finally {
updateResponseCookies()
}
}
case 'delete':
return function (names: string | string[]) {
if (Array.isArray(names)) {
names.forEach((name) => modifiedCookies.add(name))
} else {
modifiedCookies.add(names)
}
try {
return cookies.delete(names)
} finally {
updateResponseCookies()
}
}
case 'set':
return function (
...args:
| [string, string]
| [
options: NonNullable<
ReturnType<InstanceType<typeof RequestCookies>['get']>
>
]
) {
const [key, value] = args
if (typeof key === 'string') {
modifiedCookies.add(key)
try {
return cookies.set(key, value!)
} finally {
updateResponseCookies()
}
}
modifiedCookies.add(key.name)
try {
return cookies.set(key)
} finally {
updateResponseCookies()
}
}
default:
return ReflectAdapter.get(target, prop, receiver)
}
},
})
}
}

View file

@ -42,6 +42,24 @@ createNextDescribe(
const res = (await browser.elementByCss('h1').text()) || ''
return res.includes('Mozilla') ? 'UA' : ''
}, 'UA')
// Set cookies
await browser.elementByCss('#setCookie').click()
await check(async () => {
const res = (await browser.elementByCss('h1').text()) || ''
const id = res.split(':')
return id[0] === id[1] && id[0] === id[2] && id[0]
? 'same'
: 'different'
}, 'same')
})
it('should support setting cookies in route handlers with the correct overrides', async () => {
const res = await next.fetch('/handler')
const setCookieHeader = res.headers.get('set-cookie')
expect(setCookieHeader).toInclude('bar=bar2;')
expect(setCookieHeader).toInclude('baz=baz2;')
expect(setCookieHeader).toInclude('foo=foo1;')
})
it('should support formData and redirect', async () => {

View file

@ -0,0 +1,12 @@
import { cookies } from 'next/headers'
export const GET = async () => {
cookies().set('foo', 'foo1')
cookies().set('bar', 'bar1')
return new Response('Hello, world!', {
headers: [
['Set-Cookie', 'bar=bar2'],
['Set-Cookie', 'baz=baz2'],
],
})
}

View file

@ -9,3 +9,8 @@ export async function getCookie(name) {
export async function getHeader(name) {
return headers().get(name)
}
export async function setCookie(name, value) {
cookies().set(name, value)
return cookies().get(name)
}

View file

@ -1,6 +1,6 @@
import UI from './ui'
import { getCookie, getHeader } from './actions'
import { getCookie, getHeader, setCookie } from './actions'
import { validator } from './validator'
export default function Page() {
@ -9,6 +9,7 @@ export default function Page() {
<UI
getCookie={getCookie}
getHeader={getHeader}
setCookie={setCookie}
getAuthedUppercase={validator(async (str) => {
'use server'
return prefix + ' ' + str.toUpperCase()

View file

@ -2,7 +2,12 @@
import { useState } from 'react'
export default function UI({ getCookie, getHeader, getAuthedUppercase }) {
export default function UI({
getCookie,
getHeader,
setCookie,
getAuthedUppercase,
}) {
const [result, setResult] = useState('')
return (
@ -20,6 +25,23 @@ export default function UI({ getCookie, getHeader, getAuthedUppercase }) {
>
getCookie
</button>
<button
id="setCookie"
onClick={async () => {
// set cookie on server side
const random = Math.random()
const res = await setCookie('random-server', random)
setResult(
random +
':' +
res.value +
':' +
document.cookie.match(/random-server=([^;]+)/)?.[1]
)
}}
>
setCookie
</button>
<button
id="header"
onClick={async () => {