feature: InsertScriptMiddleware.js 支持编码方式 `deflate` 和 `br`，不再是只有 `gzip` 了 (#299)

packages/mitmproxy/src/lib/proxy/middleware/InsertScriptMiddleware.js

@@ -1,23 +1,47 @@
+const log = require('../../../utils/util.log')
 const through = require('through2')
 const zlib = require('zlib')
 
-const httpUtil = {}
-httpUtil.getCharset = function (res) {
-  const contentType = res.getHeader('content-type')
-  const reg = /charset=(.*)/
-  const matched = contentType.match(reg)
-  if (matched) {
-    return matched[1]
+// 编解码器
+const codecMap = {
+  gzip: {
+    createCompressor: () => zlib.createGzip(),
+    createDecompressor: () => zlib.createGunzip()
+  },
+  deflate: {
+    createCompressor: () => zlib.createDeflate(),
+    createDecompressor: () => zlib.createInflate()
+  },
+  br: {
+    createCompressor: () => zlib.createBrotliCompress(),
+    createDecompressor: () => zlib.createBrotliDecompress()
   }
-  return 'utf-8'
 }
-httpUtil.isGzip = function (res) {
-  const contentEncoding = res.headers['content-encoding']
-  return !!(contentEncoding && contentEncoding.toLowerCase() === 'gzip')
-}
-httpUtil.isHtml = function (res) {
-  const contentType = res.headers['content-type']
-  return (typeof contentType !== 'undefined') && /text\/html|application\/xhtml\+xml/.test(contentType)
+const supportedEncodings = Object.keys(codecMap)
+const supportedEncodingsStr = supportedEncodings.join(', ')
+
+const httpUtil = {
+  // 获取响应内容编码
+  getContentEncoding (res) {
+    const encoding = res.headers['content-encoding']
+    if (encoding) {
+      return encoding.toLowerCase()
+    }
+    return null
+  },
+  // 获取编解码器
+  getCodec (encoding) {
+    return codecMap[encoding]
+  },
+  // 获取支持的编解码器名称字符串
+  supportedEncodingsStr () {
+    return supportedEncodingsStr
+  },
+  // 是否HTML代码
+  isHtml (res) {
+    const contentType = res.headers['content-type']
+    return (typeof contentType !== 'undefined') && /text\/html|application\/xhtml\+xml/.test(contentType)
+  }
 }
 const HEAD = Buffer.from('</head>')
 const HEAD_UP = Buffer.from('</HEAD>')
@@ -25,16 +49,18 @@ const BODY = Buffer.from('</body>')
 const BODY_UP = Buffer.from('</BODY>')
 
 function chunkByteReplace (_this, chunk, enc, callback, append) {
-  if (append && append.head) {
-    const ret = injectScriptIntoHtml([HEAD, HEAD_UP], chunk, append.head)
-    if (ret != null) {
-      chunk = ret
+  if (append) {
+    if (append.head) {
+      const ret = injectScriptIntoHtml([HEAD, HEAD_UP], chunk, append.head)
+      if (ret != null) {
+        chunk = ret
+      }
     }
-  }
-  if (append && append.body) {
-    const ret = injectScriptIntoHtml([BODY, BODY_UP], chunk, append.body)
-    if (ret != null) {
-      chunk = ret
+    if (append.body) {
+      const ret = injectScriptIntoHtml([BODY, BODY_UP], chunk, append.body)
+      if (ret != null) {
+        chunk = ret
+      }
     }
   }
   _this.push(chunk)
@@ -57,10 +83,41 @@ function injectScriptIntoHtml (tags, chunk, script) {
   return null
 }
 
+function handleResponseHeaders (res, proxyRes) {
+  Object.keys(proxyRes.headers).forEach(function (key) {
+    if (proxyRes.headers[key] !== undefined) {
+      // let newkey = key.replace(/^[a-z]|-[a-z]/g, (match) => {
+      //   return match.toUpperCase()
+      // })
+      const newkey = key
+      if (key === 'content-length') {
+        // do nothing
+        return
+      }
+      if (key === 'content-security-policy') {
+        // content-security-policy
+        let policy = proxyRes.headers[key]
+        const reg = /script-src ([^:]*);/i
+        const matched = policy.match(reg)
+        if (matched) {
+          if (matched[1].indexOf('self') < 0) {
+            policy = policy.replace('script-src', 'script-src \'self\' ')
+          }
+        }
+        res.setHeader(newkey, policy)
+        return
+      }
+
+      res.setHeader(newkey, proxyRes.headers[key])
+    }
+  })
+
+  res.writeHead(proxyRes.statusCode)
+}
+
 const contextPath = '/____ds_script____/'
 const monkey = require('../../monkey')
 module.exports = {
-
   requestIntercept (context, req, res, ssl, next) {
     const { rOptions, log, setting } = context
     if (rOptions.path.indexOf(contextPath) !== 0) {
@@ -87,7 +144,7 @@ module.exports = {
     return true
   },
   responseInterceptor (req, res, proxyReq, proxyRes, ssl, next, append) {
-    if (!append.head && !append.body) {
+    if (append == null || (!append.head && !append.body)) {
       next()
       return
     }
@@ -99,50 +156,38 @@ module.exports = {
     if (!isHtml || contentLengthIsZero) {
       next()
       return
-    } else {
-      Object.keys(proxyRes.headers).forEach(function (key) {
-        if (proxyRes.headers[key] !== undefined) {
-          // let newkey = key.replace(/^[a-z]|-[a-z]/g, (match) => {
-          //   return match.toUpperCase()
-          // })
-          const newkey = key
-          if (isHtml && key === 'content-length') {
-            // do nothing
-            return
-          }
-          if (isHtml && key === 'content-security-policy') {
-            // content-security-policy
-            let policy = proxyRes.headers[key]
-            const reg = /script-src ([^:]*);/i
-            const matched = policy.match(reg)
-            if (matched) {
-              if (matched[1].indexOf('self') < 0) {
-                policy = policy.replace('script-src', 'script-src \'self\' ')
-              }
-            }
-            res.setHeader(newkey, policy)
-            return
-          }
-
-          res.setHeader(newkey, proxyRes.headers[key])
-        }
-      })
-
-      res.writeHead(proxyRes.statusCode)
-
-      const isGzip = httpUtil.isGzip(proxyRes)
-
-      if (isGzip) {
-        proxyRes.pipe(new zlib.Gunzip())
-          .pipe(through(function (chunk, enc, callback) {
-            chunkByteReplace(this, chunk, enc, callback, append)
-          })).pipe(new zlib.Gzip()).pipe(res)
-      } else {
-        proxyRes.pipe(through(function (chunk, enc, callback) {
-          chunkByteReplace(this, chunk, enc, callback, append)
-        })).pipe(res)
-      }
     }
+
+    // 先处理头信息
+    handleResponseHeaders(res, proxyRes)
+
+    // 获取响应内容编码
+    const encoding = httpUtil.getContentEncoding(proxyRes)
+    if (encoding) {
+      // 获取编解码器
+      const codec = httpUtil.getCodec(encoding)
+      if (codec) {
+        proxyRes
+          .pipe(codec.createDecompressor()) // 解码
+          .pipe(through(function (chunk, enc, callback) {
+            // 插入head和body
+            chunkByteReplace(this, chunk, enc, callback, append)
+          }))
+          .pipe(codec.createCompressor()) // 编码
+          .pipe(res)
+      } else {
+        log.error(`InsertScriptMiddleware.responseInterceptor(): 暂不支持编码方式 ${encoding}, 目前支持:`, httpUtil.supportedEncodingsStr())
+      }
+    } else {
+      proxyRes
+        .pipe(through(function (chunk, enc, callback) {
+          chunkByteReplace(this, chunk, enc, callback, append)
+        }))
+        .pipe(res)
+    }
+
     next()
-  }
+  },
+  httpUtil,
+  handleResponseHeaders
 }