Enable responses from custom middleware

auth0 · Jun 23, 2023 · c3ec92a · c3ec92a
1 parent 616f6f9
commit c3ec92a
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 7 deletions.
diff --git a/src/helpers/with-middleware-auth-required.ts b/src/helpers/with-middleware-auth-required.ts
@@ -82,13 +82,13 @@ export default function withMiddlewareAuthRequiredFactory(
       const res = await (middleware && middleware(...args));
 
       if (res) {
-        const headers = new Headers(res.headers);
-        const cookies = headers.get('set-cookie')?.split(', ') || [];
-        const authCookies = authRes.headers.get('set-cookie')?.split(', ') || [];
-        if (cookies.length || authCookies.length) {
-          headers.set('set-cookie', [...authCookies, ...cookies].join(', '));
+        const nextRes = new NextResponse(res.body, res);
+        for (const cookie of authRes.cookies.getAll()) {
+          if (!nextRes.cookies.get(cookie.name)) {
+            nextRes.cookies.set(cookie);
+          }
         }
-        return NextResponse.next({ ...res, status: res.status, headers });
+        return nextRes;
       } else {
         return authRes;
       }

diff --git a/tests/helpers/with-middleware-auth-required.test.ts b/tests/helpers/with-middleware-auth-required.test.ts
@@ -178,7 +178,44 @@ describe('with-middleware-auth-required', () => {
     };
     const res = await setup({ user: { name: 'dave' }, middleware });
     expect(res.status).toEqual(200);
-    expect(res.headers.get('set-cookie')).toMatch(/^appSession=.+, foo=bar;/);
+    // @ts-expect-errors ts dom doesn't have getAll
+    expect(res.headers.getAll('set-cookie')).toHaveLength(2);
+    expect(res.headers.get('set-cookie')).toMatch(/appSession=/);
+    expect(res.headers.get('set-cookie')).toMatch(/foo=bar;/);
+  });
+
+  test('should allow responses from custom middleware', async () => {
+    const middleware = () => {
+      return NextResponse.json({ foo: 'bar' });
+    };
+    const res = await setup({ user: { name: 'dave' }, middleware });
+    expect(res.status).toEqual(200);
+    await expect(res.json()).resolves.toEqual({ foo: 'bar' });
+  });
+
+  test('should allow ReadableStream responses from custom middleware', async () => {
+    const middleware = () => {
+      // @ts-expect-errors ts dom doesn't have json
+      return new Response(Response.json({ foo: 'bar' }).body);
+    };
+    const res = await setup({ user: { name: 'dave' }, middleware });
+    expect(res.status).toEqual(200);
+    await expect(res.json()).resolves.toEqual({ foo: 'bar' });
+  });
+
+  test('should allow responses and cookies from custom middleware', async () => {
+    const middleware = () => {
+      const res = NextResponse.json({ foo: 'bar' });
+      res.cookies.set('foo', 'bar');
+      return res;
+    };
+    const res = await setup({ user: { name: 'dave' }, middleware });
+    expect(res.status).toEqual(200);
+    await expect(res.json()).resolves.toEqual({ foo: 'bar' });
+    // @ts-expect-errors ts dom doesn't have getAll
+    expect(res.headers.getAll('set-cookie')).toHaveLength(2);
+    expect(res.headers.get('set-cookie')).toMatch(/appSession=/);
+    expect(res.headers.get('set-cookie')).toMatch(/foo=bar;/);
   });
 
   test('should set just a custom cookie when session is not rolling', async () => {