1  /*!
  2   * on-headers
  3   * Copyright(c) 2014 Douglas Christopher Wilson
  4   * MIT Licensed
  5   */
  6  
  7  'use strict'
  8  
  9  /**
 10   * Module exports.
 11   * @public
 12   */
 13  
 14  module.exports = onHeaders
 15  
 16  var http = require('http')
 17  
 18  // older node versions don't have appendHeader
 19  var isAppendHeaderSupported = typeof http.ServerResponse.prototype.appendHeader === 'function'
 20  var set1dArray = isAppendHeaderSupported ? set1dArrayWithAppend : set1dArrayWithSet
 21  
 22  /**
 23   * Create a replacement writeHead method.
 24   *
 25   * @param {function} prevWriteHead
 26   * @param {function} listener
 27   * @private
 28   */
 29  
 30  function createWriteHead (prevWriteHead, listener) {
 31    var fired = false
 32  
 33    // return function with core name and argument list
 34    return function writeHead (statusCode) {
 35      // set headers from arguments
 36      var args = setWriteHeadHeaders.apply(this, arguments)
 37  
 38      // fire listener
 39      if (!fired) {
 40        fired = true
 41        listener.call(this)
 42  
 43        // pass-along an updated status code
 44        if (typeof args[0] === 'number' && this.statusCode !== args[0]) {
 45          args[0] = this.statusCode
 46          args.length = 1
 47        }
 48      }
 49  
 50      return prevWriteHead.apply(this, args)
 51    }
 52  }
 53  
 54  /**
 55   * Execute a listener when a response is about to write headers.
 56   *
 57   * @param {object} res
 58   * @return {function} listener
 59   * @public
 60   */
 61  
 62  function onHeaders (res, listener) {
 63    if (!res) {
 64      throw new TypeError('argument res is required')
 65    }
 66  
 67    if (typeof listener !== 'function') {
 68      throw new TypeError('argument listener must be a function')
 69    }
 70  
 71    res.writeHead = createWriteHead(res.writeHead, listener)
 72  }
 73  
 74  /**
 75   * Set headers contained in array on the response object.
 76   *
 77   * @param {object} res
 78   * @param {array} headers
 79   * @private
 80   */
 81  
 82  function setHeadersFromArray (res, headers) {
 83    if (headers.length && Array.isArray(headers[0])) {
 84      // 2D
 85      set2dArray(res, headers)
 86    } else {
 87      // 1D
 88      if (headers.length % 2 !== 0) {
 89        throw new TypeError('headers array is malformed')
 90      }
 91  
 92      set1dArray(res, headers)
 93    }
 94  }
 95  
 96  /**
 97   * Set headers contained in object on the response object.
 98   *
 99   * @param {object} res
100   * @param {object} headers
101   * @private
102   */
103  
104  function setHeadersFromObject (res, headers) {
105    var keys = Object.keys(headers)
106    for (var i = 0; i < keys.length; i++) {
107      var k = keys[i]
108      if (k) res.setHeader(k, headers[k])
109    }
110  }
111  
112  /**
113   * Set headers and other properties on the response object.
114   *
115   * @param {number} statusCode
116   * @private
117   */
118  
119  function setWriteHeadHeaders (statusCode) {
120    var length = arguments.length
121    var headerIndex = length > 1 && typeof arguments[1] === 'string'
122      ? 2
123      : 1
124  
125    var headers = length >= headerIndex + 1
126      ? arguments[headerIndex]
127      : undefined
128  
129    this.statusCode = statusCode
130  
131    if (Array.isArray(headers)) {
132      // handle array case
133      setHeadersFromArray(this, headers)
134    } else if (headers) {
135      // handle object case
136      setHeadersFromObject(this, headers)
137    }
138  
139    // copy leading arguments
140    var args = new Array(Math.min(length, headerIndex))
141    for (var i = 0; i < args.length; i++) {
142      args[i] = arguments[i]
143    }
144  
145    return args
146  }
147  
148  function set2dArray (res, headers) {
149    var key
150    for (var i = 0; i < headers.length; i++) {
151      key = headers[i][0]
152      if (key) {
153        res.setHeader(key, headers[i][1])
154      }
155    }
156  }
157  
158  function set1dArrayWithAppend (res, headers) {
159    for (var i = 0; i < headers.length; i += 2) {
160      res.removeHeader(headers[i])
161    }
162  
163    var key
164    for (var j = 0; j < headers.length; j += 2) {
165      key = headers[j]
166      if (key) {
167        res.appendHeader(key, headers[j + 1])
168      }
169    }
170  }
171  
172  function set1dArrayWithSet (res, headers) {
173    var key
174    for (var i = 0; i < headers.length; i += 2) {
175      key = headers[i]
176      if (key) {
177        res.setHeader(key, headers[i + 1])
178      }
179    }
180  }