From 5e88852ff019bcb28bc68e93c0a57fce8063aac3 Mon Sep 17 00:00:00 2001
From: Daniele Belardi <dwon.dnl@gmail.com>
Date: Fri, 1 Oct 2021 11:17:40 +0200
Subject: [PATCH] feat: add custom undici instance

---
 README.md                      | 10 ++++-
 lib/request.js                 | 24 +++++++----
 test/custom-undici-instance.js | 77 ++++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+), 9 deletions(-)
 create mode 100644 test/custom-undici-instance.js

diff --git a/README.md b/README.md
index 6c248be9..9e4d4a89 100644
--- a/README.md
+++ b/README.md
@@ -104,9 +104,17 @@ proxy.register(require('fastify-reply-from'), {
   }
 })
 ```
-
 See undici own options for more configurations.
 
+You can also pass the plugin a custom instance:
+
+```js
+proxy.register(require('fastify-reply-from'), {
+  base: 'http://localhost:3001/',
+  undici: new undici.Pool('http://localhost:3001')
+})
+```
+
 #### `http`
 
 Set the `http` option to `true` or to an Object to use 
diff --git a/lib/request.js b/lib/request.js
index 006a1f11..1fb20c3c 100644
--- a/lib/request.js
+++ b/lib/request.js
@@ -18,9 +18,15 @@ function shouldUseUndici (opts) {
   return true
 }
 
+function isUndiciInstance (obj) {
+  return obj instanceof undici.Pool ||
+    obj instanceof undici.Client ||
+    obj instanceof undici.Dispatcher
+}
+
 function buildRequest (opts) {
   const isHttp2 = !!opts.http2
-  const isUndici = shouldUseUndici(opts)
+  const hasUndiciOptions = shouldUseUndici(opts)
   const requests = {
     'http:': http,
     'https:': https,
@@ -33,7 +39,7 @@ function buildRequest (opts) {
   const undiciOpts = opts.undici || {}
   let http2Client
   let undiciAgent
-  let undiciPool
+  let undiciInstance
   let agents
 
   if (isHttp2) {
@@ -53,12 +59,14 @@ function buildRequest (opts) {
 
   if (isHttp2) {
     return { request: handleHttp2Req, close, retryOnError: 'ECONNRESET' }
-  } else if (isUndici) {
+  } else if (hasUndiciOptions) {
     if (opts.base && opts.base.startsWith('unix+')) {
       const undiciOpts = getUndiciOptions(opts.undici)
       undiciOpts.socketPath = decodeURIComponent(new URL(opts.base).host)
       const protocol = opts.base.startsWith('unix+https') ? 'https' : 'http'
-      undiciPool = new undici.Pool(protocol + '://localhost', undiciOpts)
+      undiciInstance = new undici.Pool(protocol + '://localhost', undiciOpts)
+    } else if (isUndiciInstance(opts.undici)) {
+      undiciInstance = opts.undici
     } else {
       undiciAgent = new undici.Agent(getUndiciOptions(opts.undici))
     }
@@ -68,9 +76,9 @@ function buildRequest (opts) {
   }
 
   function close () {
-    if (isUndici) {
+    if (hasUndiciOptions) {
       undiciAgent && undiciAgent.destroy()
-      undiciPool && undiciPool.destroy()
+      undiciInstance && undiciInstance.destroy()
     } else if (!isHttp2) {
       agents['http:'].destroy()
       agents['https:'].destroy()
@@ -114,8 +122,8 @@ function buildRequest (opts) {
 
     let pool
 
-    if (undiciPool) {
-      pool = undiciPool
+    if (undiciInstance) {
+      pool = undiciInstance
     } else if (!baseUrl && opts.url.protocol.startsWith('unix')) {
       done(new Error('unix socket not supported with undici yet'))
       return
diff --git a/test/custom-undici-instance.js b/test/custom-undici-instance.js
new file mode 100644
index 00000000..00242360
--- /dev/null
+++ b/test/custom-undici-instance.js
@@ -0,0 +1,77 @@
+'use strict'
+
+const t = require('tap')
+const Fastify = require('fastify')
+const undici = require('undici')
+const http = require('http')
+const get = require('simple-get').concat
+const From = require('..')
+
+const target = http.createServer((req, res) => {
+  t.pass('request proxied')
+  t.equal(req.method, 'GET')
+  t.equal(req.url, '/')
+  t.equal(req.headers.connection, 'keep-alive')
+  res.statusCode = 205
+  res.setHeader('Content-Type', 'text/plain')
+  res.setHeader('x-my-header', 'hello!')
+  res.end('hello world')
+})
+
+t.test('use a custom instance of \'undici\'', async t => {
+  t.teardown(target.close.bind(target))
+
+  await new Promise((resolve, reject) => target.listen(0, err => err ? reject(err) : resolve()))
+
+  t.test('custom Pool', t => {
+    const instance = Fastify()
+    t.teardown(instance.close.bind(instance))
+    instance.register(From, {
+      base: `http://localhost:${target.address().port}`,
+      undici: new undici.Pool(`http://localhost:${target.address().port}`)
+    })
+
+    instance.get('/', (request, reply) => {
+      reply.from()
+    })
+
+    instance.listen(0, (err) => {
+      t.error(err)
+
+      get(`http://localhost:${instance.server.address().port}`, (err, res, data) => {
+        t.error(err)
+        t.equal(res.headers['content-type'], 'text/plain')
+        t.equal(res.headers['x-my-header'], 'hello!')
+        t.equal(res.statusCode, 205)
+        t.equal(data.toString(), 'hello world')
+        t.end()
+      })
+    })
+  })
+
+  t.test('custom Client', t => {
+    const instance = Fastify()
+    t.teardown(instance.close.bind(instance))
+    instance.register(From, {
+      base: `http://localhost:${target.address().port}`,
+      undici: new undici.Client(`http://localhost:${target.address().port}`)
+    })
+
+    instance.get('/', (request, reply) => {
+      reply.from()
+    })
+
+    instance.listen(0, (err) => {
+      t.error(err)
+
+      get(`http://localhost:${instance.server.address().port}`, (err, res, data) => {
+        t.error(err)
+        t.equal(res.headers['content-type'], 'text/plain')
+        t.equal(res.headers['x-my-header'], 'hello!')
+        t.equal(res.statusCode, 205)
+        t.equal(data.toString(), 'hello world')
+        t.end()
+      })
+    })
+  })
+})