From b31312247eb73d69dd5258f10c706f090f0ab17b Mon Sep 17 00:00:00 2001
From: Simon Schoonjans <simon.schoonjans99@gmail.com>
Date: Mon, 18 Nov 2024 17:48:29 +0100
Subject: [PATCH] fix: return 502 on invalid upstream status code (#386)

* fix: return 502 on invalid upstream status code

* chore: move warning log into error handler

Co-authored-by: Dan Castillo <dan.castillo@jasper.ai>
Signed-off-by: Simon Schoonjans <simon.schoonjans99@gmail.com>

* chore: apply suggestion to always map to badGatewayError

---------

Signed-off-by: Simon Schoonjans <simon.schoonjans99@gmail.com>
Co-authored-by: Dan Castillo <dan.castillo@jasper.ai>
---
 index.js                                  | 11 ++++-
 lib/errors.js                             |  1 +
 test/on-invalid-upstream-response.test.js | 52 +++++++++++++++++++++++
 3 files changed, 62 insertions(+), 2 deletions(-)
 create mode 100644 test/on-invalid-upstream-response.test.js

diff --git a/index.js b/index.js
index 359bf5cb..ae3d86eb 100644
--- a/index.js
+++ b/index.js
@@ -20,7 +20,8 @@ const {
   ConnectionResetError,
   ConnectTimeoutError,
   UndiciSocketError,
-  InternalServerError
+  InternalServerError,
+  BadGatewayError
 } = require('./lib/errors')
 
 const fastifyReplyFrom = fp(function from (fastify, opts, next) {
@@ -199,7 +200,13 @@ const fastifyReplyFrom = fp(function from (fastify, opts, next) {
       } else {
         copyHeaders(rewriteHeaders(res.headers, this.request), this)
       }
-      this.code(res.statusCode)
+      try {
+        this.code(res.statusCode)
+      } catch (err) {
+        // Since we know `FST_ERR_BAD_STATUS_CODE` will be recieved
+        onError(this, { error: new BadGatewayError() })
+        this.request.log.warn(err, 'response has invalid status code')
+      }
       if (onResponse) {
         onResponse(this.request, this, res)
       } else {
diff --git a/lib/errors.js b/lib/errors.js
index 87919b2e..83e27324 100644
--- a/lib/errors.js
+++ b/lib/errors.js
@@ -12,3 +12,4 @@ module.exports.ConnectionResetError = createError('ECONNRESET', 'Connection Rese
 module.exports.ConnectTimeoutError = createError('UND_ERR_CONNECT_TIMEOUT', 'Connect Timeout Error', 500)
 module.exports.UndiciSocketError = createError('UND_ERR_SOCKET', 'Undici Socket Error', 500)
 module.exports.InternalServerError = createError('FST_REPLY_FROM_INTERNAL_SERVER_ERROR', '%s', 500)
+module.exports.BadGatewayError = createError('FST_REPLY_FROM_BAD_GATEWAY', 'Bad Gateway', 502)
diff --git a/test/on-invalid-upstream-response.test.js b/test/on-invalid-upstream-response.test.js
new file mode 100644
index 00000000..994c8ddb
--- /dev/null
+++ b/test/on-invalid-upstream-response.test.js
@@ -0,0 +1,52 @@
+'use strict'
+
+const t = require('tap')
+const Fastify = require('fastify')
+const From = require('..')
+const http = require('node:http')
+const get = require('simple-get').concat
+
+const instance = Fastify()
+instance.register(From)
+
+t.plan(8)
+t.teardown(instance.close.bind(instance))
+
+const target = http.createServer((req, res) => {
+  t.pass('request proxied')
+  t.equal(req.method, 'GET')
+  res.statusCode = 888
+  res.end('non-standard status code')
+})
+
+instance.get('/', (_, reply) => {
+  reply.from(`http://localhost:${target.address().port}`, {
+    onResponse: (_, reply, res) => {
+      t.equal(res.statusCode, 888)
+    }
+  })
+})
+
+t.teardown(target.close.bind(target))
+
+instance.listen({ port: 0 }, (err) => {
+  t.error(err)
+
+  target.listen({ port: 0 }, (err) => {
+    t.error(err)
+
+    get(
+      `http://localhost:${instance.server.address().port}`,
+      (err, res, data) => {
+        t.error(err)
+        t.equal(res.statusCode, 502)
+        t.same(JSON.parse(data), {
+          statusCode: 502,
+          code: 'FST_REPLY_FROM_BAD_GATEWAY',
+          error: 'Bad Gateway',
+          message: 'Bad Gateway'
+        })
+      }
+    )
+  })
+})