From eaa1bf87dac4e7dbfa55f2dd7a14f1a89648c124 Mon Sep 17 00:00:00 2001
From: amio <amio.cn@gmail.com>
Date: Thu, 21 Nov 2019 00:35:05 +0800
Subject: [PATCH] fix: auto choose style base on hostname

---
 libs/serve-404.ts   |  7 -------
 libs/serve-badge.ts | 14 ++++++++++++--
 package-lock.json   | 13 +++++++++++++
 package.json        |  1 +
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/libs/serve-404.ts b/libs/serve-404.ts
index 3c0ac23..9d0989f 100644
--- a/libs/serve-404.ts
+++ b/libs/serve-404.ts
@@ -10,12 +10,5 @@ export default function (req, res) {
 
   const { query } = url.parse(req.url, true)
 
-  if (query.style === undefined) {
-    const host = req.headers['x-forwarded-host'] || req.headers.host
-    if (host.startsWith('flat')) {
-      query.style = 'flat'
-    }
-  }
-
   serveBadge(req, res, { code: 404, params, query: query as any })
 }
diff --git a/libs/serve-badge.ts b/libs/serve-badge.ts
index e61b144..0a3e244 100644
--- a/libs/serve-badge.ts
+++ b/libs/serve-badge.ts
@@ -1,6 +1,8 @@
 import { badgen } from 'badgen'
 import icons from 'badgen-icons'
+import originalUrl from 'original-url'
 
+import { IncomingMessage, ServerResponse } from 'http'
 import { BadgenParams } from './types'
 
 type ServeBadgeOptions = {
@@ -10,7 +12,7 @@ type ServeBadgeOptions = {
   params: BadgenParams
 }
 
-export default function (req, res, options: ServeBadgeOptions) {
+export default function (req: IncomingMessage, res: ServerResponse, options: ServeBadgeOptions) {
   const { code = 200, sMaxAge = 10800, query = {}, params } = options
 
   const { subject, status, color } = params
@@ -25,7 +27,7 @@ export default function (req, res, options: ServeBadgeOptions) {
     subject: typeof label !== 'undefined' ? label : subject,
     status: transformStatus(status, { list }),
     color: query.color || color,
-    style: (query.style || process.env.BADGE_STYLE) as 'flat' || 'classic',
+    style: autoBadgeStyle(req, query.style),
     icon: _icon.src,
     iconWidth: parseInt(iconWidth || _icon.width || '13', 10),
     scale: parseFloat(scale || '1'),
@@ -39,6 +41,14 @@ export default function (req, res, options: ServeBadgeOptions) {
   res.end(badge)
 }
 
+type BadgeStyle = 'flat' | undefined
+
+function autoBadgeStyle (req: IncomingMessage, fallbackValue?: string): BadgeStyle {
+  const isFlat = (fallbackValue || process.env.BADGE_STYLE === 'flat')
+    || originalUrl(req).hostname.includes('flat')
+  return isFlat ? 'flat' : undefined
+}
+
 function transformStatus (status: any, { list }) {
   status = String(status)
 
diff --git a/package-lock.json b/package-lock.json
index 567ceae..770f3de 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -4242,6 +4242,11 @@
       "integrity": "sha512-Fs9VRguL0gqGHkXS5GQiMCr1VhZBxz0JnJs4JmMp/2jL18Fmbzvv7vOFRU+U8TBkHEE/CX1qDXzJplVULgsLeg==",
       "dev": true
     },
+    "forwarded-parse": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmjs.org/forwarded-parse/-/forwarded-parse-2.1.0.tgz",
+      "integrity": "sha512-as9a7Xelt0CvdUy7/qxrY73dZq2vMx49F556fwjjFrUyzq5uHHfeLgD2cCq/6P4ZvusGZzjD6aL2NdgGdS5Cew=="
+    },
     "fragment-cache": {
       "version": "0.2.1",
       "resolved": "https://registry.npmjs.org/fragment-cache/-/fragment-cache-0.2.1.tgz",
@@ -6841,6 +6846,14 @@
         "wcwidth": "^1.0.1"
       }
     },
+    "original-url": {
+      "version": "1.2.3",
+      "resolved": "https://registry.npmjs.org/original-url/-/original-url-1.2.3.tgz",
+      "integrity": "sha512-BYm+pKYLtS4mVe/mgT3YKGtWV5HzN/XKiaIu1aK4rsxyjuHeTW9N+xVBEpJcY1onB3nccfH0RbzUEoimMqFUHQ==",
+      "requires": {
+        "forwarded-parse": "^2.1.0"
+      }
+    },
     "os-browserify": {
       "version": "0.3.0",
       "resolved": "https://registry.npmjs.org/os-browserify/-/os-browserify-0.3.0.tgz",
diff --git a/package.json b/package.json
index ecf4dba..0c790e0 100644
--- a/package.json
+++ b/package.json
@@ -35,6 +35,7 @@
     "micro": "^9.3.4",
     "millify": "^3.1.0",
     "my-way": "^1.0.2",
+    "original-url": "^1.2.3",
     "react-debounce-render": "^5.0.0",
     "semver": "^6.3.0",
     "serve-handler": "^6.1.2",