Rewrite location header (#25)

Author: mcollina

index.js

@@ -14,11 +14,28 @@ module.exports = async function (fastify, opts) {
   fromOpts.base = opts.upstream
   fromOpts.prefix = undefined
 
+  const oldRewriteHeaders = (opts.replyOptions || {}).rewriteHeaders
+  const replyOpts = Object.assign({}, opts.replyOpts, {
+    rewriteHeaders
+  })
+  fromOpts.rewriteHeaders = rewriteHeaders
+
   fastify.register(From, fromOpts)
 
   fastify.addContentTypeParser('application/json', bodyParser)
   fastify.addContentTypeParser('*', bodyParser)
 
+  function rewriteHeaders (headers) {
+    const location = headers.location
+    if (location) {
+      headers.location = location.replace(rewritePrefix, fastify.basePath)
+    }
+    if (oldRewriteHeaders) {
+      headers = oldRewriteHeaders(headers)
+    }
+    return headers
+  }
+
   function bodyParser (req, done) {
     done(null, req)
   }
@@ -29,6 +46,6 @@ module.exports = async function (fastify, opts) {
   function reply (request, reply) {
     var dest = request.req.url
     dest = dest.replace(this.basePath, rewritePrefix)
-    reply.from(dest || '/', opts.replyOptions)
+    reply.from(dest || '/', replyOpts)
   }
 }

test.js

@@ -18,6 +18,7 @@ async function run () {
 
   origin.post('/this-has-data', async (request, reply) => {
     if (request.body.hello === 'world') {
+      reply.header('location', '/something')
       return { something: 'posted' }
     }
     throw new Error('kaboom')
@@ -184,7 +185,7 @@ async function run () {
     t.match(headers, { 'x-test': 'test' })
   })
 
-  test('keepPrefix true', async (t) => {
+  test('rewritePrefix', async (t) => {
     const proxyServer = Fastify()
 
     proxyServer.register(proxy, {
@@ -202,6 +203,27 @@ async function run () {
     const firstProxyPrefix = await got(`http://localhost:${proxyServer.server.address().port}/api/a`)
     t.equal(firstProxyPrefix.body, 'this is /api2/a')
   })
+
+  test('rewrite location headers', async (t) => {
+    const proxyServer = Fastify()
+
+    proxyServer.register(proxy, {
+      upstream: `http://localhost:${origin.server.address().port}`,
+      prefix: '/api'
+    })
+
+    await proxyServer.listen(0)
+
+    t.tearDown(() => {
+      proxyServer.close()
+    })
+
+    const { headers: { location } } = await got(`http://localhost:${proxyServer.server.address().port}/api/this-has-data`, {
+      body: { hello: 'world' },
+      json: true
+    })
+    t.equal(location, '/api/something')
+  })
 }
 
 run()