on-headers: 一个用于在返回请求头数据前执行回调函数的工具库

本文是“源码解读”系列第 2 篇。

概述

on-headers 是一个 Node.js 端的工具库，用来在所有请求头数据设置之后、发送给客户端之前，为你提供一个回调函数用来执行。

下面是它的一个使用案例：

var http = require('http')
var onHeaders = require('on-headers')

http
  .createServer(onRequest)
  .listen(3000)

function addPoweredBy () {
  // set if not set by end of request
  if (!this.getHeader('X-Powered-By')) {
    this.setHeader('X-Powered-By', 'Node.js')
  }
}

function onRequest (req, res) {
  onHeaders(res, addPoweredBy)

  res.setHeader('Content-Type', 'text/plain')
  res.end('hello!')
}

on-headers 的使用场景一般是在响应头发送给客户端之前做一些数据检查。

按照作者的话说^[1]，这个库是为了模拟 res 上的 headers 事件：res.on('headers', listener)，目前 Node.js 原生还不支持。

实现思路

知道 on-headers 包的作用后，我们再来了解一下它的实现思路：通过复写 response.writeHead()^[2] 方法。

认识 response.writeHead() 方法

与 res.setHeader() 的异同点

response.writeHead() 方法的作用类似 res.setHeader()，不过还需要我们指定 HTTP 状态码。

拿下面的例子^[3]来说：

const http = require('node:http');
const hostname = '127.0.0.1';
const port = 3000;
const server = http.createServer((req, res) => {
  res.statusCode = 200;
  res.setHeader('Content-Type', 'text/plain');
  res.end('Hello Worldn');
});
server.listen(port, hostname, () => {
  console.log(`Server running at http://${hostname}:${port}/`);
});

换成 response.writeHead() 实现是这样的：

const http = require('node:http');
const hostname = '127.0.0.1';
const port = 3000;
const server = http.createServer((req, res) => {
  res.writeHead(200, {'Content-Type': 'text/plain' })
  res.end('Hello Worldn');
});
server.listen(port, hostname, () => {
  console.log(`Server running at http://${hostname}:${port}/`);
});

除了需要设置 HTTP 状态码，response.writeHead() 与 response.setHeader() 的不同之处还包括。

一、res.writeHead() 可以同时设置多个响应头

res.writeHead(200, {
  'X-Powered-By': 'Node.js',
  'Content-Type': 'text/plain',
})

二、res.writeHead() 之后就不能再使用 res.setHeader() 了，会报错

res.writeHead(200, {'Content-Type': 'text/plain'});
res.setHeader('X-Foo', 'bar');

// Error [ERR_HTTP_HEADERS_SENT]: Cannot set headers after they are sent to the client

这是因为在调用 res.writeHead() 之后，请求头已经发送会客户端了。

response.writeHead() 总是会被调用

response.writeHead() 方法是作为最终返回响应头的操作。即使你没有显式调用，它也会在你调用 res.end() 方法的时候检查，确保 response.writeHead() 方法被调用过了。

我举一个例子：

// server.js
import http from 'node:http'

http
  .createServer(onRequest)
  .listen(3000)

function onRequest (req, res) {
  const old = res.writeHead
  res.writeHead = function (...args) {
    console.log(';)', args)
    old.call(res, ...args)
  }

  res.statusCode = 200
  res.setHeader('Content-Type', 'text/html');
  res.end('ok');
}

这里们复写了 res.writeHead() 的实现，仅仅做了一个 log 打印操作。node server.js 启动后，浏览器访问 http://localhost:3000/，可以看见在终端上看到打印结果。

$ node .index.js
;) [ 200 ]

分析源码

在理解了 on-headers 的实现原理和 res.writeHead() 方法的作用后，再来看一下 on-headers 源码。

❝

我们基于最新的 v1.0.2^[4] 版本讲解。

❞

package.json 里没有指定 "main" 字段，就知道入口源码是在项目根目录下的 index.js 中。由于 on-headers 的实现思路不复杂，因此代码都写在 index.js 里面了，连注释一共才 132 行。

先来看默认导出。

/**
 * Module exports.
 * @public
 */

module.exports = onHeaders

再看一下 onHeaders() 函数。

onHeaders()

onHeaders() 函数接收 res、listener 两次参数。res 就是响应对象，listener 就是我们要注册的回调函数了。

/**
 * Execute a listener when a response is about to write headers.
 *
 * @param {object} res
 * @return {function} listener
 * @public
 */

function onHeaders (res, listener) {
  if (!res) {
    throw new TypeError('argument res is required')
  }

  if (typeof listener !== 'function') {
    throw new TypeError('argument listener must be a function')
  }

  res.writeHead = createWriteHead(res.writeHead, listener)
}

这里先检查了参数是否合规，然后使用 createWriteHead() 的返回值复写了 res.writeHead() 函数。

接下来，再来看 createWriteHead() 的实现。

createWriteHead()

createWriteHead() 函数接收原始版本的 res.writeHead() 方法和监听函数 listener。

/**
 * Create a replacement writeHead method.
 *
 * @param {function} prevWriteHead
 * @param {function} listener
 * @private
 */

function createWriteHead (prevWriteHead, listener) {
  // return function with core name and argument list
  return function writeHead (statusCode) {
    // ...
  }
}

这里返回的 writeHead (statusCode) {} 就是我们的复写版本了。在这里我们会在调用原始版本的 res.writeHead() 方法前，执行监听函数。

调用 listener

首先添加调用 listener 的逻辑。

function createWriteHead (prevWriteHead, listener) {
  var fired = false
  
  // return function with core name and argument list
  return function writeHead (statusCode) {
    // fire listener
    if (!fired) {
      fired = true
      listener.call(this)
    }

    // ...
  }
}

调用 listener 时将 response（即这里的 this）作为它的上下文对象。同时，为了保证监听函数只被唯一调用一次，我们引入了标识变量 fired。

当然，我们的 on-headers listener 执行时机是在所有请求头数据设置之后、发送给客户端之前。因此，我们还需要调用在 listener 之前，设置好响应头。

  // return function with core name and argument list
  return function writeHead (statusCode) {
   // set headers from arguments
    var args = setWriteHeadHeaders.apply(this, arguments)
    
    // fire listener
    if (!fired) {
      fired = true
      listener.call(this)
    }

    // ...
  }

这里额外引入了一个 setWriteHeadHeaders() 函数，我们来看一下。

setWriteHeadHeaders()

setWriteHeadHeaders() 函数的作用就是将 writeHead() 接收到的 headers 数据，通过 res.setHeader() 设置，这样在调用 listener 时，就能确保所有请求头数据都设置好了。

/**
 * Set headers and other properties on the response object.
 *
 * @param {number} statusCode
 * @private
 */

function setWriteHeadHeaders (statusCode) {
  var length = arguments.length
  var headerIndex = length > 1 && typeof arguments[1] === 'string'
    ? 2
    : 1

  var headers = length >= headerIndex + 1
    ? arguments[headerIndex]
    : undefined

  this.statusCode = statusCode

  if (Array.isArray(headers)) {
    // handle array case
    setHeadersFromArray(this, headers)
  } else if (headers) {
    // handle object case
    setHeadersFromObject(this, headers)
  }

  // copy leading arguments
  var args = new Array(Math.min(length, headerIndex))
  for (var i = 0; i < args.length; i++) {
    args[i] = arguments[i]
  }

  return args
}

setWriteHeadHeaders() 的实现不是很复杂，主要是做了 4 件事。

1. 找到 headers 参数。因为 writeHead 函数签名^[5]是 response.writeHead(statusCode[, statusMessage][, headers])，就是说第二个可能是表示 statusMessage 的字符串，那这个时候 headers 参数应该是在第 3 个位置（对应索引 2）
2. 通过 res.setHeader() 设置请求头数据。由于 headers 参数既可能是对象（常用）也可能是数组（不常用），因此对应分别实现了 setHeadersFromArray() 和 setHeadersFromObject() 两个工具函数。实现比较简单，我们直接贴出代码——就是循环遍历每个 Header 的 key 和 value，代入 res.setHeader() 中即可

/**
 * Set headers contained in array on the response object.
 *
 * @param {object} res
 * @param {array} headers
 * @private
 */

function setHeadersFromArray (res, headers) {
  for (var i = 0; i < headers.length; i++) {
    res.setHeader(headers[i][0], headers[i][1])
  }
}

/**
 * Set headers contained in object on the response object.
 *
 * @param {object} res
 * @param {object} headers
 * @private
 */

function setHeadersFromObject (res, headers) {
  var keys = Object.keys(headers)
  for (var i = 0; i < keys.length; i++) {
    var k = keys[i]
    if (k) res.setHeader(k, headers[k])
  }
}

3. 设置 res.statusCode。由于在调用 .writeHead() 之前，外部可能已经手动设置 res.statusCode 了，以 .writeHead() 接收到的为准
4. 将 headers 之前的参数返回。比如，我们是以 [200, {'Content-Type': 'text/plain'}] 调用的，那么返回的 args 就是 [200]

现在回到 createWriteHead() 再来看一下。

回到 createWriteHead()

我们将调用 setWriteHeadHeaders() 返回的结果 args，在调用 listener 之后，透传给原始的 res.writeHead() 方法。

  // return function with core name and argument list
  return function writeHead (statusCode) {
   // set headers from arguments
    var args = setWriteHeadHeaders.apply(this, arguments)
    
    // fire listener
    if (!fired) {
      fired = true
      listener.call(this)
    }

    // 将调用 setWriteHeadHeaders() 返回的结果透传给原始 res.writeHead 方法
    return prevWriteHead.apply(this, args)
  }

这就完成了 on-headers 的所有代码实现。

总结

在完成了 on-headers 包的实现之后，我们就能知道下面的 2 块代码是等价的。

function onRequest (req, res) {
  onHeaders(res, addPoweredBy)
  res.writeHead(200, { 'Content-Type': 'text/plain' })
  res.end('hello!')
}

function onRequest (req, res) {
  res.setHeader('Content-Type', 'text/plain')
  addPoweredBy(res)
  res.writeHead(200)
  res.end('hello!')
}

知其然，知其所以然之后，我们就能更有信心的去使用这个工具库啦。

希望本篇文章能对你有所帮助。感谢阅读，Happy Coding！

参考资料