index.js 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. /*!
  2. * on-headers
  3. * Copyright(c) 2014 Douglas Christopher Wilson
  4. * MIT Licensed
  5. */
  6. 'use strict'
  7. /**
  8. * Module exports.
  9. * @public
  10. */
  11. module.exports = onHeaders
  12. /**
  13. * Create a replacement writeHead method.
  14. *
  15. * @param {function} prevWriteHead
  16. * @param {function} listener
  17. * @private
  18. */
  19. function createWriteHead (prevWriteHead, listener) {
  20. var fired = false
  21. // return function with core name and argument list
  22. return function writeHead (statusCode) {
  23. // set headers from arguments
  24. var args = setWriteHeadHeaders.apply(this, arguments)
  25. // fire listener
  26. if (!fired) {
  27. fired = true
  28. listener.call(this)
  29. // pass-along an updated status code
  30. if (typeof args[0] === 'number' && this.statusCode !== args[0]) {
  31. args[0] = this.statusCode
  32. args.length = 1
  33. }
  34. }
  35. return prevWriteHead.apply(this, args)
  36. }
  37. }
  38. /**
  39. * Execute a listener when a response is about to write headers.
  40. *
  41. * @param {object} res
  42. * @return {function} listener
  43. * @public
  44. */
  45. function onHeaders (res, listener) {
  46. if (!res) {
  47. throw new TypeError('argument res is required')
  48. }
  49. if (typeof listener !== 'function') {
  50. throw new TypeError('argument listener must be a function')
  51. }
  52. res.writeHead = createWriteHead(res.writeHead, listener)
  53. }
  54. /**
  55. * Set headers contained in array on the response object.
  56. *
  57. * @param {object} res
  58. * @param {array} headers
  59. * @private
  60. */
  61. function setHeadersFromArray (res, headers) {
  62. for (var i = 0; i < headers.length; i++) {
  63. res.setHeader(headers[i][0], headers[i][1])
  64. }
  65. }
  66. /**
  67. * Set headers contained in object on the response object.
  68. *
  69. * @param {object} res
  70. * @param {object} headers
  71. * @private
  72. */
  73. function setHeadersFromObject (res, headers) {
  74. var keys = Object.keys(headers)
  75. for (var i = 0; i < keys.length; i++) {
  76. var k = keys[i]
  77. if (k) res.setHeader(k, headers[k])
  78. }
  79. }
  80. /**
  81. * Set headers and other properties on the response object.
  82. *
  83. * @param {number} statusCode
  84. * @private
  85. */
  86. function setWriteHeadHeaders (statusCode) {
  87. var length = arguments.length
  88. var headerIndex = length > 1 && typeof arguments[1] === 'string'
  89. ? 2
  90. : 1
  91. var headers = length >= headerIndex + 1
  92. ? arguments[headerIndex]
  93. : undefined
  94. this.statusCode = statusCode
  95. if (Array.isArray(headers)) {
  96. // handle array case
  97. setHeadersFromArray(this, headers)
  98. } else if (headers) {
  99. // handle object case
  100. setHeadersFromObject(this, headers)
  101. }
  102. // copy leading arguments
  103. var args = new Array(Math.min(length, headerIndex))
  104. for (var i = 0; i < args.length; i++) {
  105. args[i] = arguments[i]
  106. }
  107. return args
  108. }