From 54e9bdc7c03b77fe748e82cac7ec0e3ff1d56aea Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=BCrg=C3=BCn=20Day=C4=B1o=C4=9Flu?=
 <gurgun.dayioglu@icloud.com>
Date: Tue, 1 Aug 2023 17:50:04 +0300
Subject: [PATCH] Check if onSend was ran and keep resetting the map (#242)

* set cookie map to null after sendHeaders run

* don'r reassign reply to null

* remove kReplySetCookiesHookRan from onReqHandler

* nit: fix Symbol description
---
 plugin.js           | 15 ++++++---------
 test/cookie.test.js | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/plugin.js b/plugin.js
index c321d16..b4f84c0 100644
--- a/plugin.js
+++ b/plugin.js
@@ -6,13 +6,9 @@ const cookie = require('cookie')
 const { Signer, sign, unsign } = require('./signer')
 
 const kReplySetCookies = Symbol('fastify.reply.setCookies')
+const kReplySetCookiesHookRan = Symbol('fastify.reply.setCookiesHookRan')
 
 function fastifyCookieSetCookie (reply, name, value, options) {
-  let sendHeaders = false
-  if (reply[kReplySetCookies] === null) {
-    sendHeaders = true
-    reply[kReplySetCookies] = new Map()
-  }
   const opts = Object.assign({}, options)
 
   if (opts.expires && Number.isInteger(opts.expires)) {
@@ -34,8 +30,9 @@ function fastifyCookieSetCookie (reply, name, value, options) {
 
   reply[kReplySetCookies].set(`${name};${opts.domain};${opts.path || '/'}`, { name, value, opts })
 
-  if (sendHeaders) {
+  if (reply[kReplySetCookiesHookRan]) {
     setCookies(reply)
+    reply[kReplySetCookies].clear()
   }
 
   return reply
@@ -103,9 +100,8 @@ function fastifyCookieOnSendHandler (fastifyReq, fastifyRes, payload, done) {
     setCookies(fastifyRes)
   }
 
-  // Explicitly set the property to null so that we can
-  // check if the header was already set
-  fastifyRes[kReplySetCookies] = null
+  fastifyRes[kReplySetCookies].clear()
+  fastifyRes[kReplySetCookiesHookRan] = true
 
   done()
 }
@@ -147,6 +143,7 @@ function plugin (fastify, options, next) {
 
   fastify.decorateRequest('cookies', null)
   fastify.decorateReply(kReplySetCookies, null)
+  fastify.decorateReply(kReplySetCookiesHookRan, false)
 
   fastify.decorateReply('cookie', setCookie)
   fastify.decorateReply('setCookie', setCookie)
diff --git a/test/cookie.test.js b/test/cookie.test.js
index d09ed77..0dacac1 100644
--- a/test/cookie.test.js
+++ b/test/cookie.test.js
@@ -1187,3 +1187,42 @@ test('cookies get set correctly if set inside onSend', (t) => {
     t.equal(cookies[0].path, '/')
   })
 })
+
+test('cookies get set correctly if set inside multiple onSends', (t) => {
+  t.plan(10)
+  const fastify = Fastify()
+  fastify.register(plugin)
+
+  fastify.addHook('onSend', async (req, reply, payload) => {
+    reply.setCookie('foo', 'foo', { path: '/' })
+  })
+
+  fastify.addHook('onSend', async (req, reply, payload) => {
+    reply.setCookie('foo', 'foos', { path: '/' })
+    return payload
+  })
+
+  fastify.get('/test1', (req, reply) => {
+    reply
+      .send({ hello: 'world' })
+  })
+
+  fastify.inject({
+    method: 'GET',
+    url: '/test1'
+  }, (err, res) => {
+    t.error(err)
+    t.equal(res.statusCode, 200)
+    t.same(JSON.parse(res.body), { hello: 'world' })
+
+    const cookies = res.cookies
+    t.equal(cookies.length, 2)
+    t.equal(cookies[0].name, 'foo')
+    t.equal(cookies[0].value, 'foo')
+    t.equal(cookies[0].path, '/')
+
+    t.equal(cookies[1].name, 'foo')
+    t.equal(cookies[1].value, 'foos')
+    t.equal(cookies[1].path, '/')
+  })
+})