From 212bbbda764676195d2ccaf1ee60f35daaa24076 Mon Sep 17 00:00:00 2001
From: Vlad Sirenko <sirenkovladd@gmail.com>
Date: Sun, 3 Sep 2023 03:03:15 +0300
Subject: [PATCH] parsing socket

---
 index.js              |   6 +-
 lib/do-inject.js      |  27 +++--
 lib/request.js        | 165 ++++++++++++++-------------
 lib/response.js       | 193 +++++++++++---------------------
 lib/socket.js         | 138 +++++++++++++++++++++++
 test/index.test.js    | 251 +++++++++++++++++++++++++++++++-----------
 test/response.test.js |   4 +-
 types/index.d.ts      |   5 +
 types/index.test-d.ts |   5 +-
 9 files changed, 511 insertions(+), 283 deletions(-)
 create mode 100644 lib/socket.js

diff --git a/index.js b/index.js
index ee8ec15..c640b17 100644
--- a/index.js
+++ b/index.js
@@ -1,7 +1,7 @@
 'use strict'
 
 const Request = require('./lib/request')
-const Response = require('./lib/response')
+const { Response } = require('./lib/response')
 const Chain = require('./lib/chain')
 const doInject = require('./lib/do-inject')
 
@@ -25,3 +25,7 @@ module.exports = inject
 module.exports.default = inject
 module.exports.inject = inject
 module.exports.isInjection = isInjection
+module.exports.errors = {
+  ...Request.errors,
+  ...Response.errors
+}
diff --git a/lib/do-inject.js b/lib/do-inject.js
index 6a17279..0b5f065 100644
--- a/lib/do-inject.js
+++ b/lib/do-inject.js
@@ -3,11 +3,12 @@
 const assert = require('assert')
 const optsValidator = require('./configValidator')
 const Request = require('./request')
-const Response = require('./response')
+const { Response, once } = require('./response')
+const { Readable, addAbortSignal } = require('stream')
 
 function promisify (fn) {
   if (fn) {
-    return { ret: Promise.resolve(), cb: fn }
+    return { ret: Promise.resolve(), cb: once(fn) }
   }
   let resolve, reject
   const ret = new Promise((_resolve, _reject) => {
@@ -23,15 +24,11 @@ function promisify (fn) {
 }
 
 function makeRequest (dispatchFunc, server, req, res) {
-  req.once('error', function (err) {
-    if (this.destroyed) res.destroy(err)
+  req.socket.once('close', function () {
+    res.emit('close')
   })
 
-  req.once('close', function () {
-    if (this.destroyed && !this._error) res.destroy()
-  })
-
-  return req.prepare(() => dispatchFunc.call(server, req, res))
+  return req.prepare(() => dispatchFunc.call(server, req, res), (err) => { res.emit('error', err) })
 }
 
 function doInject (dispatchFunc, options, callback) {
@@ -62,6 +59,18 @@ function doInject (dispatchFunc, options, callback) {
   const req = new RequestConstructor(options)
   const res = new Response(req, cb)
 
+  if (options.signal) {
+    const r = new Readable()
+    r.once('error', (err) => {
+      cb(err)
+      res.destroy(err)
+    })
+    res.once('close', () => {
+      r.destroy()
+    })
+    addAbortSignal(options.signal, r)
+  }
+
   return Promise.resolve().then(() => makeRequest(dispatchFunc, server, req, res)).then(() => ret)
 }
 
diff --git a/lib/request.js b/lib/request.js
index aeaf474..3f12efc 100644
--- a/lib/request.js
+++ b/lib/request.js
@@ -2,13 +2,14 @@
 
 /* eslint no-prototype-builtins: 0 */
 
-const { Readable, addAbortSignal } = require('stream')
+const { Readable } = require('stream')
 const cookie = require('cookie')
 const assert = require('assert')
 const warning = require('process-warning')()
 
 const parseURL = require('./parseURL')
-const { EventEmitter } = require('events')
+const { IncomingMessage } = require('http')
+const { Socket } = require('net')
 
 // request.connectin deprecation https://nodejs.org/api/http.html#http_request_connection
 warning.create('FastifyDeprecationLightMyRequest', 'FST_LIGHTMYREQUEST_DEP01', 'You are accessing "request.connection", use "request.socket" instead.')
@@ -31,14 +32,27 @@ function hostHeaderFromURL (parsedURL) {
  * @constructor
  * @param {String} remoteAddress the fake address to show consumers of the socket
  */
-class MockSocket extends EventEmitter {
+class MockSocket extends Socket {
   constructor (remoteAddress) {
     super()
-    this.remoteAddress = remoteAddress
+    Object.defineProperty(this, 'remoteAddress', {
+      __proto__: null,
+      configurable: false,
+      enumerable: true,
+      get: () => remoteAddress
+    })
   }
 }
 
-class Request extends Readable {
+class Request extends IncomingMessage {
+  static errors = {
+    ContentLength: class ContentLengthError extends Error {
+      constructor () {
+        super('Content length is different than the value specified by the Content-Length header')
+      }
+    }
+  }
+
   /**
    * Request
    *
@@ -49,10 +63,11 @@ class Request extends Readable {
    * @param {Object} [options.cookies]
    * @param {Object} [options.headers]
    * @param {Object} [options.query]
+   * @param {{end: boolean,split: boolean,error: boolean,close: boolean}} [options.simulate]
    * @param {any} [options.payload]
    */
   constructor (options) {
-    super({ autoDestroy: false })
+    super(new MockSocket(options.remoteAddress || '127.0.0.1'))
     const parsedURL = parseURL(options.url || options.path, options.query)
 
     this.url = parsedURL.pathname + parsedURL.search
@@ -97,8 +112,6 @@ class Request extends Readable {
       this.headers.cookie = cookieValues.join('; ')
     }
 
-    this.socket = new MockSocket(options.remoteAddress || '127.0.0.1')
-
     Object.defineProperty(this, 'connection', {
       get () {
         warning.emit('FST_LIGHTMYREQUEST_DEP01')
@@ -119,97 +132,93 @@ class Request extends Readable {
       }
     }
 
-    // Set the content-length for the corresponding payload if none set
-    if (payload && !payloadResume && !Object.prototype.hasOwnProperty.call(this.headers, 'content-length')) {
-      this.headers['content-length'] = (Buffer.isBuffer(payload) ? payload.length : Buffer.byteLength(payload)).toString()
-    }
-
     for (const header of Object.keys(this.headers)) {
       this.rawHeaders.push(header, this.headers[header])
     }
 
+    if (options.simulate?.end === false) {
+      const prevPayload = payload
+      if (payloadResume) {
+        payload = new Readable({
+          read (n) {
+            prevPayload.read(n)
+          }
+        })
+        prevPayload.on('data', (d) => {
+          payload.push(d)
+        })
+      } else {
+        payload = new Readable({
+          read (n) {
+            if (prevPayload) this.push(prevPayload)
+            this.pause()
+          }
+        })
+      }
+    }
+
     // Use _lightMyRequest namespace to avoid collision with Node
     this._lightMyRequest = {
       payload,
       isDone: false,
       simulate: options.simulate || {}
     }
-
-    const signal = options.signal
-    /* istanbul ignore if  */
-    if (signal) {
-      addAbortSignal(signal, this)
-    }
   }
 
-  prepare (next) {
-    const payload = this._lightMyRequest.payload
-    if (!payload || typeof payload.resume !== 'function') { // does not quack like a stream
-      return next()
+  getLength (payload) {
+    if (typeof payload === 'string') {
+      return Buffer.byteLength(payload)
     }
 
-    const chunks = []
-
-    payload.on('data', (chunk) => chunks.push(Buffer.from(chunk)))
-
-    payload.on('end', () => {
-      const payload = Buffer.concat(chunks)
-      this.headers['content-length'] = this.headers['content-length'] || ('' + payload.length)
-      this._lightMyRequest.payload = payload
-      return next()
-    })
-
-    // Force to resume the stream. Needed for Stream 1
-    payload.resume()
+    return payload.length
   }
 
-  _read (size) {
-    setImmediate(() => {
-      if (this._lightMyRequest.isDone) {
-        // 'end' defaults to true
-        if (this._lightMyRequest.simulate.end !== false) {
-          this.push(null)
-        }
-
-        return
-      }
-
-      this._lightMyRequest.isDone = true
-
-      if (this._lightMyRequest.payload) {
-        if (this._lightMyRequest.simulate.split) {
-          this.push(this._lightMyRequest.payload.slice(0, 1))
-          this.push(this._lightMyRequest.payload.slice(1))
+  prepare (next, onError) {
+    let payload = this._lightMyRequest.payload
+    this.complete = true
+    if (payload) {
+      if (typeof payload.resume !== 'function') {
+        const length = this.getLength(payload)
+        if (this.headers['content-length']) {
+          if (this.headers['content-length'].toString() > length.toString()) {
+            return onError(new Request.errors.ContentLength())
+          }
+          payload = payload.slice(0, this.headers['content-length'])
         } else {
-          this.push(this._lightMyRequest.payload)
+          this.headers['content-length'] = length?.toString()
         }
-      }
-
-      if (this._lightMyRequest.simulate.error) {
-        this.emit('error', new Error('Simulated'))
-      }
-
-      if (this._lightMyRequest.simulate.close) {
-        this.emit('close')
-      }
-
-      // 'end' defaults to true
-      if (this._lightMyRequest.simulate.end !== false) {
+        this.push(payload)
         this.push(null)
+      } else {
+        let i = 0
+        const max = this.headers['content-length'] ? parseInt(this.headers['content-length'], 10) : null
+        payload.on('data', (chunk) => {
+          if (max != null) {
+            if (max > i && max <= i + chunk.length) {
+              this.push(chunk.slice(0, max - i))
+            }
+          } else {
+            this.push(chunk)
+          }
+          i += chunk.length
+        })
+        payload.on('end', () => {
+          if (max != null) {
+            if (max > i) {
+              return onError(new Request.errors.ContentLength())
+            }
+          }
+          this.push(null)
+        })
+        payload.resume()
       }
-    })
-  }
-
-  destroy (error) {
-    if (this.destroyed || this._lightMyRequest.isDone) return
-    this.destroyed = true
-
-    if (error) {
-      this._error = true
-      process.nextTick(() => this.emit('error', error))
+    } else {
+      if (this.headers['content-length'] && this.headers['content-length'] !== '0') {
+        return onError(new Request.errors.ContentLength())
+      }
+      this.push(null)
     }
-
-    process.nextTick(() => this.emit('close'))
+    return next()
   }
 }
 
diff --git a/lib/response.js b/lib/response.js
index 4a6945a..d9be629 100644
--- a/lib/response.js
+++ b/lib/response.js
@@ -1,48 +1,77 @@
 'use strict'
 
 const http = require('http')
-const { Writable } = require('stream')
 
 const setCookie = require('set-cookie-parser')
+const MySocket = require('./socket')
+
+function once (cb) {
+  let called = false
+  return function () {
+    if (called) return
+    called = true
+    cb.apply(this, arguments)
+  }
+}
+
+module.exports.once = once
 
 // Throws away all written data to prevent response from buffering payload
-function getNullSocket () {
-  return new Writable({
-    write (chunk, encoding, callback) {
-      setImmediate(callback)
-    }
-  })
-}
 
 class Response extends http.ServerResponse {
+  static errors = {
+    SocketHangUpError: class SocketHangUpError extends Error {
+      constructor () {
+        super('socket hang up')
+        this.code = 'ECONNRESET'
+      }
+    }
+  }
+
+  /**
+   * @param {import('./request').Request} req
+   * @param {(err: Error, data: any) => void} onEnd
+   * @param {http.Server} server
+   */
   constructor (req, onEnd) {
     super(req)
-    this._lightMyRequest = { headers: null, trailers: {}, payloadChunks: [] }
+    onEnd = once(onEnd)
+    this._lightSocket = new MySocket()
     this.setHeader('foo', 'bar'); this.removeHeader('foo')
-    this.assignSocket(getNullSocket())
-    let called = false
-    const onEndSuccess = (payload) => {
-      // no need to early-return if already called because this handler is bound `once`
-      called = true
-      process.nextTick(() => onEnd(null, payload))
-    }
+    this.assignSocket(this._lightSocket)
 
-    const onEndFailure = (err) => {
-      if (called) return
-      called = true
-      process.nextTick(() => onEnd(err, null))
+    const onEndCb = (err) => {
+      if (err) {
+        return process.nextTick(() => onEnd(err))
+      }
+      const res = this.generatePayload(req)
+      if (res.end) {
+        return process.nextTick(() => onEnd(null, res))
+      }
+      process.nextTick(() => onEnd(new Response.errors.SocketHangUpError()))
     }
 
     this.once('finish', () => {
-      const res = this.generatePayload(req)
-      onEndSuccess(res)
+      this.destroyed = true
+      this._closed = true
+      this.emit('close')
     })
 
-    this.socket.once('error', onEndFailure)
+    this.once('close', () => {
+      onEndCb()
+    })
+
+    this.socket.once('error', () => {
+      onEndCb(new Response.errors.SocketHangUpError())
+    })
 
-    this.once('error', onEndFailure)
+    this.socket.once('close', () => {
+      process.nextTick(() => onEndCb())
+    })
 
-    this.once('close', onEndFailure)
+    this.once('error', (err) => {
+      onEndCb(err)
+    })
   }
 
   setTimeout (msecs, callback) {
@@ -53,55 +82,6 @@ class Response extends http.ServerResponse {
     return this
   }
 
-  writeHead (...opt) {
-    const result = super.writeHead(...opt)
-
-    this.copyHeaders()
-
-    return result
-  }
-
-  write (data, encoding, callback) {
-    if (this.timeoutHandle) {
-      clearTimeout(this.timeoutHandle)
-    }
-    super.write(data, encoding, callback)
-    this._lightMyRequest.payloadChunks.push(Buffer.from(data, encoding))
-    return true
-  }
-
-  end (data, encoding, callback) {
-    if (data) {
-      this.write(data, encoding)
-    }
-
-    super.end(callback)
-
-    this.emit('finish')
-
-    // We need to emit 'close' otherwise stream.finished() would
-    // not pick it up on Node v16
-
-    this.destroy()
-  }
-
-  destroy (error) {
-    if (this.destroyed) return
-    this.destroyed = true
-
-    if (error) {
-      process.nextTick(() => this.emit('error', error))
-    }
-
-    process.nextTick(() => this.emit('close'))
-  }
-
-  addTrailers (trailers) {
-    for (const key in trailers) {
-      this._lightMyRequest.trailers[key.toLowerCase().trim()] = trailers[key].toString().trim()
-    }
-  }
-
   /**
    * @private
    * @param {Request} req
@@ -109,73 +89,32 @@ class Response extends http.ServerResponse {
    */
   generatePayload (req) {
     // This seems only to happen when using `fastify-express` - see https://github.com/fastify/fastify-express/issues/47
-    /* istanbul ignore if */
-    if (this._lightMyRequest.headers === null) {
-      this.copyHeaders()
-    }
-    this.serializeHeaders()
     // Prepare response object
+    const state = this._lightSocket.getState()
+    const body = state.body.toString()
     const res = {
       raw: {
         res: this,
         req
       },
-      headers: this._lightMyRequest.headers,
+      headers: state.headers,
       statusCode: this.statusCode,
       statusMessage: this.statusMessage,
-      trailers: {},
+      trailers: state.trailers,
+      rawPayload: state.body,
+      end: state.isEnd,
+      payload: body,
+      body,
+      json: function parseJsonPayload () {
+        return JSON.parse(res.payload)
+      },
       get cookies () {
         return setCookie.parse(this)
       }
     }
 
-    // Prepare payload and trailers
-    const rawBuffer = Buffer.concat(this._lightMyRequest.payloadChunks)
-    res.rawPayload = rawBuffer
-
-    // we keep both of them for compatibility reasons
-    res.payload = rawBuffer.toString()
-    res.body = res.payload
-    res.trailers = this._lightMyRequest.trailers
-
-    // Prepare payload parsers
-    res.json = function parseJsonPayload () {
-      return JSON.parse(res.payload)
-    }
-
     return res
   }
-
-  /**
-   * @private
-   */
-  serializeHeaders () {
-    const headers = this._lightMyRequest.headers
-
-    for (const headerName of Object.keys(headers)) {
-      const headerValue = headers[headerName]
-      if (Array.isArray(headerValue)) {
-        headers[headerName] = headerValue.map(value => '' + value)
-      } else {
-        headers[headerName] = '' + headerValue
-      }
-    }
-  }
-
-  /**
-   * @private
-   */
-  copyHeaders () {
-    this._lightMyRequest.headers = Object.assign({}, this.getHeaders())
-
-    // Add raw headers
-    ;['Date', 'Connection', 'Transfer-Encoding'].forEach((name) => {
-      const value = this.getHeader(name)
-      if (value) {
-        this._lightMyRequest.headers[name.toLowerCase()] = value
-      }
-    })
-  }
 }
 
-module.exports = Response
+module.exports.Response = Response
diff --git a/lib/socket.js b/lib/socket.js
new file mode 100644
index 0000000..28e522d
--- /dev/null
+++ b/lib/socket.js
@@ -0,0 +1,138 @@
+const { Socket } = require('net')
+
+const crlfBuf = Buffer.from('\r\n')
+
+class State {
+  constructor () {
+    this.state = 'firstHead'
+    this.headers = {}
+    this.trailers = {}
+    this.body = []
+    this.waitSizeBody = 0
+    this.rawBody = []
+  }
+
+  /**
+   * @param {Uint8Array | string} chunk
+   * @param {BufferEncoding} [encoding]
+   */
+  write (chunk, encoding) {
+    if (typeof chunk === 'string') {
+      chunk = Buffer.from(chunk, encoding)
+    }
+    this.rawBody.push(chunk)
+    this.process(Buffer.concat(this.rawBody))
+  }
+
+  /**
+   * @private
+   * @param {Buffer} buffer
+   * @returns
+   */
+  process (buffer) {
+    if (!buffer.length) {
+      this.rawBody = []
+      return
+    };
+    if (this.state === 'body') {
+      if (buffer.length < this.waitSizeBody) {
+        this.body.push(buffer)
+        this.waitSizeBody -= buffer.length
+        this.rawBody = []
+        return
+      }
+      this.body.push(buffer.subarray(0, this.waitSizeBody))
+      const size = this.waitSizeBody
+      this.waitSizeBody = 0
+      this.state = 'afterBody'
+      this.process(buffer.subarray(size))
+      return
+    }
+    this.rawBody = [buffer]
+    const i = buffer.indexOf(crlfBuf)
+    if (i === -1) {
+      return
+    };
+    if (this.state === 'firstHead') {
+      this.state = 'head'
+      this.process(buffer.subarray(i + 2))
+      return
+    }
+    if (this.state === 'head') {
+      const line = buffer.subarray(0, i).toString()
+      if (line) {
+        const [, key, value] = line.match(/^([^:]+): (.*)$/)
+        if (this.headers[key.toLowerCase()]) {
+          if (!Array.isArray(this.headers[key.toLowerCase()])) {
+            this.headers[key.toLowerCase()] = [this.headers[key.toLowerCase()]]
+          }
+          this.headers[key.toLowerCase()].push(value)
+        } else {
+          this.headers[key.toLowerCase()] = value
+        }
+      } else if (this.headers['content-length']) {
+        this.waitSizeBody = parseInt(this.headers['content-length'])
+        if (this.waitSizeBody) {
+          this.state = 'body'
+        } else {
+          this.state = 'trailers'
+        }
+      } else {
+        this.state = 'bodySize'
+      }
+      this.process(buffer.subarray(i + 2))
+      return
+    }
+    if (this.state === 'bodySize') {
+      const chunk = buffer.subarray(0, i).toString()
+      this.waitSizeBody = parseInt(chunk.toString(), 16)
+      if (this.waitSizeBody !== 0) {
+        this.state = 'body'
+      } else {
+        this.state = 'trailers'
+      }
+      this.process(buffer.subarray(i + 2))
+      return
+    }
+    if (this.state === 'afterBody') {
+      this.state = 'bodySize'
+      this.process(buffer.subarray(i + 2))
+    }
+    if (this.state === 'trailers') {
+      const line = buffer.subarray(0, i).toString()
+      if (line) {
+        const [, key, value] = line.match(/^([^:]+): (.*)$/)
+        this.trailers[key.toLowerCase()] = value
+      } else {
+        this.state = 'end'
+      }
+      this.process(buffer.subarray(i + 2))
+    }
+  }
+}
+
+class MySocket extends Socket {
+  _myData = new State()
+
+  /**
+   * @param {Uint8Array | string} chunk
+   * @param {BufferEncoding | (err?: Error) => void} [encoding]
+   * @param {(err?: Error) => void} [callback]
+   */
+  write (chunk, encoding, callback) {
+    this._myData.write(chunk, encoding)
+    callback && setImmediate(callback)
+    return true
+  }
+
+  getState () {
+    return {
+      headers: this._myData.headers,
+      body: Buffer.concat(this._myData.body),
+      trailers: this._myData.trailers,
+      isEnd: ['end', 'trailers', 'afterBody', 'bodySize'].includes(this._myData.state)
+    }
+  }
+}
+
+module.exports = MySocket
diff --git a/test/index.test.js b/test/index.test.js
index 8633501..9d1aa6e 100644
--- a/test/index.test.js
+++ b/test/index.test.js
@@ -27,8 +27,40 @@ const httpMethods = [
   'trace'
 ]
 
+const parseQuery = url => {
+  const parsedURL = parseURL(url)
+  return qs.parse(parsedURL.search.slice(1))
+}
+
+function getTestStream (encoding) {
+  const word = 'hi'
+  let i = 0
+
+  const stream = new Readable({
+    read (n) {
+      this.push(word[i] ? word[i++] : null)
+    }
+  })
+
+  if (encoding) {
+    stream.setEncoding(encoding)
+  }
+
+  return stream
+}
+
+function readStream (stream, callback) {
+  const chunks = []
+
+  stream.on('data', (chunk) => chunks.push(chunk))
+
+  stream.on('end', () => {
+    return callback(Buffer.concat(chunks))
+  })
+}
+
 test('returns non-chunked payload', (t) => {
-  t.plan(6)
+  t.plan(7)
   const output = 'example.com:8080|/hello'
 
   const dispatch = function (req, res) {
@@ -42,7 +74,10 @@ test('returns non-chunked payload', (t) => {
     t.error(err)
     t.equal(res.statusCode, 200)
     t.equal(res.statusMessage, 'Super')
+    t.ok(res.headers.date)
     t.strictSame(res.headers, {
+      date: res.headers.date,
+      connection: 'keep-alive',
       'x-extra': 'hello',
       'content-type': 'text/plain',
       'content-length': output.length.toString()
@@ -53,7 +88,7 @@ test('returns non-chunked payload', (t) => {
 })
 
 test('returns single buffer payload', (t) => {
-  t.plan(3)
+  t.plan(6)
   const dispatch = function (req, res) {
     res.writeHead(200, { 'Content-Type': 'text/plain' })
     res.end(req.headers.host + '|' + req.url)
@@ -61,6 +96,9 @@ test('returns single buffer payload', (t) => {
 
   inject(dispatch, { url: 'http://example.com:8080/hello' }, (err, res) => {
     t.error(err)
+    t.ok(res.headers.date)
+    t.ok(res.headers.connection)
+    t.equal(res.headers['transfer-encoding'], 'chunked')
     t.equal(res.payload, 'example.com:8080|/hello')
     t.equal(res.rawPayload.toString(), 'example.com:8080|/hello')
   })
@@ -156,12 +194,13 @@ test('passes a socket which emits events like a normal one does', (t) => {
   const dispatch = function (req, res) {
     res.writeHead(200, { 'Content-Type': 'text/plain' })
     req.socket.on('timeout', () => {})
+    res.write('some test ')
     res.end('added')
   }
 
   inject(dispatch, { method: 'GET', url: 'http://example.com:8080/hello' }, (err, res) => {
     t.error(err)
-    t.equal(res.payload, 'added')
+    t.equal(res.payload, 'some test added')
   })
 })
 
@@ -191,11 +230,6 @@ test('includes deprecated connection on request', (t) => {
   })
 })
 
-const parseQuery = url => {
-  const parsedURL = parseURL(url)
-  return qs.parse(parsedURL.search.slice(1))
-}
-
 test('passes query', (t) => {
   t.plan(2)
 
@@ -362,7 +396,7 @@ test('includes default https port in host header', (t) => {
 })
 
 test('optionally accepts an object as url', (t) => {
-  t.plan(3)
+  t.plan(5)
   const output = 'example.com:8080|/hello?test=1234'
 
   const dispatch = function (req, res) {
@@ -382,6 +416,8 @@ test('optionally accepts an object as url', (t) => {
 
   inject(dispatch, { url }, (err, res) => {
     t.error(err)
+    t.ok(res.headers.date)
+    t.ok(res.headers.connection)
     t.notOk(res.headers['transfer-encoding'])
     t.equal(res.payload, output)
   })
@@ -401,7 +437,7 @@ test('leaves user-agent unmodified', (t) => {
 })
 
 test('returns chunked payload', (t) => {
-  t.plan(2)
+  t.plan(5)
   const dispatch = function (req, res) {
     res.writeHead(200, 'OK')
     res.write('a')
@@ -411,6 +447,9 @@ test('returns chunked payload', (t) => {
 
   inject(dispatch, { method: 'GET', url: '/' }, (err, res) => {
     t.error(err)
+    t.ok(res.headers.date)
+    t.ok(res.headers.connection)
+    t.equal(res.headers['transfer-encoding'], 'chunked')
     t.equal(res.payload, 'ab')
   })
 })
@@ -662,13 +701,12 @@ test('adds a content-length header if none set when payload specified', (t) => {
 test('retains a content-length header when payload specified', (t) => {
   t.plan(2)
   const dispatch = function (req, res) {
-    res.writeHead(200, { 'Content-Type': 'text/plain' })
-    res.end(req.headers['content-length'])
+    t.fail()
   }
 
-  inject(dispatch, { method: 'POST', url: '/test', payload: '', headers: { 'content-length': '10' } }, (err, res) => {
-    t.error(err)
-    t.equal(res.payload, '10')
+  inject(dispatch, { method: 'POST', url: '/test', payload: '1', headers: { 'content-length': '10' } }, (err, res) => {
+    t.equal(err instanceof inject.errors.ContentLength, true)
+    t.error(res)
   })
 })
 
@@ -712,8 +750,8 @@ test('can override stream payload content-length header', (t) => {
   const headers = { 'content-length': '100' }
 
   inject(dispatch, { method: 'POST', url: '/', payload: getTestStream(), headers }, (err, res) => {
-    t.error(err)
-    t.equal(res.payload, '100')
+    t.equal(err instanceof inject.errors.ContentLength, true)
+    t.error(res)
   })
 })
 
@@ -721,7 +759,7 @@ test('can override stream payload content-length header without request content-
   t.plan(1)
   const dispatch = function (req, res) {
     res.writeHead(200, { 'Content-Type': 'text/plain' })
-    t.equal(req.headers['content-length'], '2')
+    t.error(req.headers['content-length'])
   }
 
   inject(dispatch, { method: 'POST', url: '/', payload: getTestStream() }, () => {})
@@ -757,7 +795,6 @@ test('_read() plays payload', (t) => {
     })
 
     req.on('end', () => {
-      res.writeHead(200, { 'Content-Length': 0 })
       res.end(buffer)
       req.destroy()
     })
@@ -782,7 +819,6 @@ test('simulates split', (t) => {
     })
 
     req.on('end', () => {
-      res.writeHead(200, { 'Content-Length': 0 })
       res.end(buffer)
       req.destroy()
     })
@@ -795,14 +831,14 @@ test('simulates split', (t) => {
   })
 })
 
-test('simulates error', (t) => {
+t.skip('simulates error', (t) => {
   t.plan(2)
   const dispatch = function (req, res) {
     req.on('readable', () => {
     })
 
     req.on('error', () => {
-      res.writeHead(200, { 'Content-Length': 0 })
+      res.writeHead(200, { 'Content-Length': 5 })
       res.end('error')
     })
   }
@@ -865,7 +901,7 @@ test('simulates close', (t) => {
     })
 
     req.on('close', () => {
-      res.writeHead(200, { 'Content-Length': 0 })
+      res.writeHead(200, { 'Content-Length': 5 })
       res.end('close')
     })
 
@@ -1160,7 +1196,7 @@ test('chainable api: http methods should work correctly', (t) => {
     inject(dispatch)[method]('http://example.com:8080/hello')
       .end((err, res) => {
         t.error(err)
-        t.equal(res.body, method.toUpperCase())
+        t.equal(res.body, method === 'head' ? '' : method.toUpperCase())
       })
   })
 })
@@ -1494,33 +1530,6 @@ test('disabling autostart', (t) => {
   })
 })
 
-function getTestStream (encoding) {
-  const word = 'hi'
-  let i = 0
-
-  const stream = new Readable({
-    read (n) {
-      this.push(word[i] ? word[i++] : null)
-    }
-  })
-
-  if (encoding) {
-    stream.setEncoding(encoding)
-  }
-
-  return stream
-}
-
-function readStream (stream, callback) {
-  const chunks = []
-
-  stream.on('data', (chunk) => chunks.push(chunk))
-
-  stream.on('end', () => {
-    return callback(Buffer.concat(chunks))
-  })
-}
-
 test('send cookie', (t) => {
   t.plan(3)
   const dispatch = function (req, res) {
@@ -1629,6 +1638,8 @@ test('correctly handles no string headers', (t) => {
       date: date.toString(),
       true: 'true',
       false: 'false',
+      connection: 'keep-alive',
+      'transfer-encoding': 'chunked',
       'content-type': 'application/json'
     })
 
@@ -1726,7 +1737,7 @@ test('no error for response destory', (t) => {
   }
 
   inject(dispatch, { method: 'GET', url: '/' }, (err, res) => {
-    t.error(err)
+    t.equal(err instanceof inject.errors.SocketHangUpError, true)
   })
 })
 
@@ -1738,8 +1749,8 @@ test('request destory without error', (t) => {
   }
 
   inject(dispatch, { method: 'GET', url: '/' }, (err, res) => {
-    t.error(err)
-    t.equal(res, null)
+    t.equal(err instanceof inject.errors.SocketHangUpError, true)
+    t.error(res)
   })
 })
 
@@ -1753,8 +1764,8 @@ test('request destory with error', (t) => {
   }
 
   inject(dispatch, { method: 'GET', url: '/' }, (err, res) => {
-    t.equal(err, fakeError)
-    t.equal(res, null)
+    t.equal(err instanceof inject.errors.SocketHangUpError, true)
+    t.error(res)
   })
 })
 
@@ -1770,8 +1781,8 @@ test('compatible with stream.finished', (t) => {
   }
 
   inject(dispatch, { method: 'GET', url: '/' }, (err, res) => {
-    t.error(err)
-    t.equal(res, null)
+    t.equal(err instanceof inject.errors.SocketHangUpError, true)
+    t.error(res)
   })
 })
 
@@ -1787,8 +1798,8 @@ test('compatible with eos', (t) => {
   }
 
   inject(dispatch, { method: 'GET', url: '/' }, (err, res) => {
-    t.error(err)
-    t.equal(res, null)
+    t.equal(err instanceof inject.errors.SocketHangUpError, true)
+    t.error(res)
   })
 })
 
@@ -1821,15 +1832,15 @@ test('compatible with eos, passes error correctly', (t) => {
 
   const dispatch = function (req, res) {
     eos(res, (err) => {
-      t.equal(err, fakeError)
+      t.equal(err.message, 'premature close')
     })
 
     req.destroy(fakeError)
   }
 
   inject(dispatch, { method: 'GET', url: '/' }, (err, res) => {
-    t.equal(err, fakeError)
-    t.equal(res, null)
+    t.equal(err instanceof inject.errors.SocketHangUpError, true)
+    t.error(res)
   })
 })
 
@@ -1842,8 +1853,8 @@ test('multiple calls to req.destroy should not be called', (t) => {
   }
 
   inject(dispatch, { method: 'GET', url: '/' }, (err, res) => {
-    t.equal(err)
-    t.equal(res, null)
+    t.equal(err instanceof inject.errors.SocketHangUpError, true)
+    t.error(res)
   })
 })
 
@@ -2014,3 +2025,113 @@ test('request that is destroyed does not error', (t) => {
     t.equal(res.payload, 'hi')
   })
 })
+
+test('status 204', (t) => {
+  t.plan(5)
+  const dispatch = function (req, res) {
+    res.writeHead(204)
+    res.end('hello')
+  }
+
+  inject(dispatch, { method: 'POST', url: '/' }, (err, res) => {
+    t.error(err)
+    t.equal(res.payload, '')
+    t.equal(Object.keys(res.headers).length, 2)
+    t.hasProp(res.headers, 'date')
+    t.equal(res.headers.connection, 'keep-alive')
+  })
+})
+
+test('partly send body', (t) => {
+  t.plan(2)
+  const dispatch = function (req, res) {
+    res.writeHead(200, { 'content-length': 10 })
+    res.write('')
+    res.write('1234')
+    res.write('123456')
+    res.end()
+  }
+
+  inject(dispatch, { method: 'POST', url: '/' }, (err, res) => {
+    t.error(err)
+    t.equal(res.payload, '1234123456')
+  })
+})
+
+test('retains a content-length header without payload', (t) => {
+  t.plan(2)
+  const dispatch = function (req, res) {
+    t.fail()
+  }
+
+  inject(dispatch, { method: 'POST', url: '/test', headers: { 'content-length': '10' } }, (err, res) => {
+    t.equal(err instanceof inject.errors.ContentLength, true)
+    t.error(res)
+  })
+})
+
+test('content-length correct with payload', (t) => {
+  t.plan(2)
+  const dispatch = function (req, res) {
+    res.end(req.headers['content-length'])
+  }
+
+  inject(dispatch, { method: 'POST', url: '/test', payload: '1234', headers: { 'content-length': '4' } }, (err, res) => {
+    t.error(err)
+    t.equal(res.payload, '4')
+  })
+})
+
+test('content-length slice if less', (t) => {
+  t.plan(2)
+  const dispatch = function (req, res) {
+    const chunks = []
+    req.on('data', (chunk) => chunks.push(chunk))
+    req.on('end', () => {
+      res.end(Buffer.concat(chunks).toString())
+    })
+  }
+
+  inject(dispatch, { method: 'POST', url: '/test', payload: '123456', headers: { 'content-length': '4' } }, (err, res) => {
+    t.error(err)
+    t.equal(res.payload, '1234')
+  })
+})
+
+test('content-length slice if less readable', (t) => {
+  t.plan(2)
+  const dispatch = function (req, res) {
+    readStream(req, (buff) => {
+      res.writeHead(200, { 'Content-Type': 'text/plain' })
+      res.end(buff)
+    })
+  }
+
+  const payload = getTestStream()
+
+  inject(dispatch, { method: 'POST', url: '/', payload, headers: { 'content-length': '1' } }, (err, res) => {
+    t.error(err)
+    t.equal(res.payload, 'h')
+  })
+})
+
+test('simulates no end with payload Readable', (t) => {
+  t.plan(2)
+  let end = false
+  const dispatch = function (req, res) {
+    req.resume()
+    req.on('end', () => {
+      end = true
+    })
+  }
+
+  let replied = false
+  inject(dispatch, { method: 'GET', url: '/', payload: getTestStream(), simulate: { end: false } }, (notHandledErr, res) => {
+    replied = true
+  })
+
+  setTimeout(() => {
+    t.equal(end, false)
+    t.equal(replied, false)
+  }, 10)
+})
diff --git a/test/response.test.js b/test/response.test.js
index 4e1f50c..24ab362 100644
--- a/test/response.test.js
+++ b/test/response.test.js
@@ -1,13 +1,13 @@
 const { test } = require('tap')
 
-const Response = require('../lib/response')
+const { Response } = require('../lib/response')
 
 test('multiple calls to res.destroy should not be called', (t) => {
   t.plan(1)
 
   const mockReq = {}
   const res = new Response(mockReq, (err, response) => {
-    t.error(err)
+    t.equal(err instanceof Response.errors.SocketHangUpError, true)
   })
 
   res.destroy()
diff --git a/types/index.d.ts b/types/index.d.ts
index 4abe499..a8147b9 100644
--- a/types/index.d.ts
+++ b/types/index.d.ts
@@ -24,6 +24,11 @@ declare namespace inject {
     readonly aborted: boolean;
   }
 
+  export const errors: {
+    ContentLength: typeof Error,
+    SocketHangUpError: typeof Error,
+  }
+
   export interface InjectOptions {
     url?: string | {
       pathname: string
diff --git a/types/index.test-d.ts b/types/index.test-d.ts
index 4ff3b20..c6baaa0 100644
--- a/types/index.test-d.ts
+++ b/types/index.test-d.ts
@@ -1,5 +1,5 @@
 import * as http from 'http'
-import { inject, isInjection, Response, DispatchFunc, InjectOptions, Chain } from '..'
+import { inject, isInjection, Response, DispatchFunc, InjectOptions, Chain, errors } from '..'
 import { expectType, expectAssignable, expectNotAssignable } from 'tsd'
 
 expectAssignable<InjectOptions>({ url: '/' })
@@ -15,6 +15,9 @@ const dispatch: http.RequestListener = function (req, res) {
   res.end(reply)
 }
 
+expectType<typeof Error>(errors.ContentLength)
+expectType<typeof Error>(errors.SocketHangUpError)
+
 const expectResponse = function (res: Response | undefined) {
   if (!res) {
     return;