Fix @fastify/session (#240)

Signed-off-by: Matteo Collina <hello@matteocollina.com>

plugin.js

@@ -8,6 +8,11 @@ const { Signer, sign, unsign } = require('./signer')
 const kReplySetCookies = Symbol('fastify.reply.setCookies')
 
 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)) {
@@ -29,6 +34,10 @@ function fastifyCookieSetCookie (reply, name, value, options) {
 
   reply[kReplySetCookies].set(`${name};${opts.domain};${opts.path || '/'}`, { name, value, opts })
 
+  if (sendHeaders) {
+    setCookies(reply)
+  }
+
   return reply
 }
 
@@ -63,33 +72,41 @@ function onReqHandlerWrapper (fastify, hook) {
     }
 }
 
-function fastifyCookieOnSendHandler (fastifyReq, fastifyRes, payload, done) {
-  if (fastifyRes[kReplySetCookies].size) {
-    let setCookie = fastifyRes.getHeader('Set-Cookie')
-
-    /* istanbul ignore else */
-    if (setCookie === undefined) {
-      if (fastifyRes[kReplySetCookies].size === 1) {
-        for (const c of fastifyRes[kReplySetCookies].values()) {
-          fastifyRes.header('Set-Cookie', cookie.serialize(c.name, c.value, c.opts))
-        }
+function setCookies (reply) {
+  let setCookie = reply.getHeader('Set-Cookie')
 
-        return done()
+  /* istanbul ignore else */
+  if (setCookie === undefined) {
+    if (reply[kReplySetCookies].size === 1) {
+      for (const c of reply[kReplySetCookies].values()) {
+        reply.header('Set-Cookie', cookie.serialize(c.name, c.value, c.opts))
       }
 
-      setCookie = []
-    } else if (typeof setCookie === 'string') {
-      setCookie = [setCookie]
+      return
     }
 
-    for (const c of fastifyRes[kReplySetCookies].values()) {
-      setCookie.push(cookie.serialize(c.name, c.value, c.opts))
-    }
+    setCookie = []
+  } else if (typeof setCookie === 'string') {
+    setCookie = [setCookie]
+  }
 
-    fastifyRes.removeHeader('Set-Cookie')
-    fastifyRes.header('Set-Cookie', setCookie)
+  for (const c of reply[kReplySetCookies].values()) {
+    setCookie.push(cookie.serialize(c.name, c.value, c.opts))
   }
 
+  reply.removeHeader('Set-Cookie')
+  reply.header('Set-Cookie', setCookie)
+}
+
+function fastifyCookieOnSendHandler (fastifyReq, fastifyRes, payload, done) {
+  if (fastifyRes[kReplySetCookies].size) {
+    setCookies(fastifyRes)
+  }
+
+  // Explicitly set the property to null so that we can
+  // check if the header was already set
+  fastifyRes[kReplySetCookies] = null
+
   done()
 }
 

test/cookie.test.js

@@ -1156,3 +1156,34 @@ test('should update a cookie value when setCookie is called multiple times (non-
     t.ok(new Date(cookies[1].expires) < new Date())
   })
 })
+
+test('cookies get set correctly if set inside onSend', (t) => {
+  t.plan(7)
+  const fastify = Fastify()
+  fastify.register(plugin)
+
+  fastify.addHook('onSend', async (req, reply, payload) => {
+    reply.setCookie('foo', 'foo', { 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, 1)
+    t.equal(cookies[0].name, 'foo')
+    t.equal(cookies[0].value, 'foo')
+    t.equal(cookies[0].path, '/')
+  })
+})