3 files changed, 587 insertions, 0 deletions
diff --git a/lib/handler/DecoratorHandler.js b/lib/handler/DecoratorHandler.js
new file mode 100644
index 0000000..9d70a76
--- /dev/null
+++ b/lib/handler/DecoratorHandler.js
@@ -0,0 +1,35 @@
+'use strict'
+
+module.exports = class DecoratorHandler {
+  constructor (handler) {
+    this.handler = handler
+  }
+
+  onConnect (...args) {
+    return this.handler.onConnect(...args)
+  }
+
+  onError (...args) {
+    return this.handler.onError(...args)
+  }
+
+  onUpgrade (...args) {
+    return this.handler.onUpgrade(...args)
+  }
+
+  onHeaders (...args) {
+    return this.handler.onHeaders(...args)
+  }
+
+  onData (...args) {
+    return this.handler.onData(...args)
+  }
+
+  onComplete (...args) {
+    return this.handler.onComplete(...args)
+  }
+
+  onBodySent (...args) {
+    return this.handler.onBodySent(...args)
+  }
+}
diff --git a/lib/handler/RedirectHandler.js b/lib/handler/RedirectHandler.js
new file mode 100644
index 0000000..baca27e
--- /dev/null
+++ b/lib/handler/RedirectHandler.js
@@ -0,0 +1,216 @@
+'use strict'
+
+const util = require('../core/util')
+const { kBodyUsed } = require('../core/symbols')
+const assert = require('assert')
+const { InvalidArgumentError } = require('../core/errors')
+const EE = require('events')
+
+const redirectableStatusCodes = [300, 301, 302, 303, 307, 308]
+
+const kBody = Symbol('body')
+
+class BodyAsyncIterable {
+  constructor (body) {
+    this[kBody] = body
+    this[kBodyUsed] = false
+  }
+
+  async * [Symbol.asyncIterator] () {
+    assert(!this[kBodyUsed], 'disturbed')
+    this[kBodyUsed] = true
+    yield * this[kBody]
+  }
+}
+
+class RedirectHandler {
+  constructor (dispatch, maxRedirections, opts, handler) {
+    if (maxRedirections != null && (!Number.isInteger(maxRedirections) || maxRedirections < 0)) {
+      throw new InvalidArgumentError('maxRedirections must be a positive number')
+    }
+
+    util.validateHandler(handler, opts.method, opts.upgrade)
+
+    this.dispatch = dispatch
+    this.location = null
+    this.abort = null
+    this.opts = { ...opts, maxRedirections: 0 } // opts must be a copy
+    this.maxRedirections = maxRedirections
+    this.handler = handler
+    this.history = []
+
+    if (util.isStream(this.opts.body)) {
+      // TODO (fix): Provide some way for the user to cache the file to e.g. /tmp
+      // so that it can be dispatched again?
+      // TODO (fix): Do we need 100-expect support to provide a way to do this properly?
+      if (util.bodyLength(this.opts.body) === 0) {
+        this.opts.body
+          .on('data', function () {
+            assert(false)
+          })
+      }
+
+      if (typeof this.opts.body.readableDidRead !== 'boolean') {
+        this.opts.body[kBodyUsed] = false
+        EE.prototype.on.call(this.opts.body, 'data', function () {
+          this[kBodyUsed] = true
+        })
+      }
+    } else if (this.opts.body && typeof this.opts.body.pipeTo === 'function') {
+      // TODO (fix): We can't access ReadableStream internal state
+      // to determine whether or not it has been disturbed. This is just
+      // a workaround.
+      this.opts.body = new BodyAsyncIterable(this.opts.body)
+    } else if (
+      this.opts.body &&
+      typeof this.opts.body !== 'string' &&
+      !ArrayBuffer.isView(this.opts.body) &&
+      util.isIterable(this.opts.body)
+    ) {
+      // TODO: Should we allow re-using iterable if !this.opts.idempotent
+      // or through some other flag?
+      this.opts.body = new BodyAsyncIterable(this.opts.body)
+    }
+  }
+
+  onConnect (abort) {
+    this.abort = abort
+    this.handler.onConnect(abort, { history: this.history })
+  }
+
+  onUpgrade (statusCode, headers, socket) {
+    this.handler.onUpgrade(statusCode, headers, socket)
+  }
+
+  onError (error) {
+    this.handler.onError(error)
+  }
+
+  onHeaders (statusCode, headers, resume, statusText) {
+    this.location = this.history.length >= this.maxRedirections || util.isDisturbed(this.opts.body)
+      ? null
+      : parseLocation(statusCode, headers)
+
+    if (this.opts.origin) {
+      this.history.push(new URL(this.opts.path, this.opts.origin))
+    }
+
+    if (!this.location) {
+      return this.handler.onHeaders(statusCode, headers, resume, statusText)
+    }
+
+    const { origin, pathname, search } = util.parseURL(new URL(this.location, this.opts.origin && new URL(this.opts.path, this.opts.origin)))
+    const path = search ? `${pathname}${search}` : pathname
+
+    // Remove headers referring to the original URL.
+    // By default it is Host only, unless it's a 303 (see below), which removes also all Content-* headers.
+    // https://tools.ietf.org/html/rfc7231#section-6.4
+    this.opts.headers = cleanRequestHeaders(this.opts.headers, statusCode === 303, this.opts.origin !== origin)
+    this.opts.path = path
+    this.opts.origin = origin
+    this.opts.maxRedirections = 0
+    this.opts.query = null
+
+    // https://tools.ietf.org/html/rfc7231#section-6.4.4
+    // In case of HTTP 303, always replace method to be either HEAD or GET
+    if (statusCode === 303 && this.opts.method !== 'HEAD') {
+      this.opts.method = 'GET'
+      this.opts.body = null
+    }
+  }
+
+  onData (chunk) {
+    if (this.location) {
+      /*
+        https://tools.ietf.org/html/rfc7231#section-6.4
+
+        TLDR: undici always ignores 3xx response bodies.
+
+        Redirection is used to serve the requested resource from another URL, so it is assumes that
+        no body is generated (and thus can be ignored). Even though generating a body is not prohibited.
+
+        For status 301, 302, 303, 307 and 308 (the latter from RFC 7238), the specs mention that the body usually
+        (which means it's optional and not mandated) contain just an hyperlink to the value of
+        the Location response header, so the body can be ignored safely.
+
+        For status 300, which is "Multiple Choices", the spec mentions both generating a Location
+        response header AND a response body with the other possible location to follow.
+        Since the spec explicitily chooses not to specify a format for such body and leave it to
+        servers and browsers implementors, we ignore the body as there is no specified way to eventually parse it.
+      */
+    } else {
+      return this.handler.onData(chunk)
+    }
+  }
+
+  onComplete (trailers) {
+    if (this.location) {
+      /*
+        https://tools.ietf.org/html/rfc7231#section-6.4
+
+        TLDR: undici always ignores 3xx response trailers as they are not expected in case of redirections
+        and neither are useful if present.
+
+        See comment on onData method above for more detailed informations.
+      */
+
+      this.location = null
+      this.abort = null
+
+      this.dispatch(this.opts, this)
+    } else {
+      this.handler.onComplete(trailers)
+    }
+  }
+
+  onBodySent (chunk) {
+    if (this.handler.onBodySent) {
+      this.handler.onBodySent(chunk)
+    }
+  }
+}
+
+function parseLocation (statusCode, headers) {
+  if (redirectableStatusCodes.indexOf(statusCode) === -1) {
+    return null
+  }
+
+  for (let i = 0; i < headers.length; i += 2) {
+    if (headers[i].toString().toLowerCase() === 'location') {
+      return headers[i + 1]
+    }
+  }
+}
+
+// https://tools.ietf.org/html/rfc7231#section-6.4.4
+function shouldRemoveHeader (header, removeContent, unknownOrigin) {
+  return (
+    (header.length === 4 && header.toString().toLowerCase() === 'host') ||
+    (removeContent && header.toString().toLowerCase().indexOf('content-') === 0) ||
+    (unknownOrigin && header.length === 13 && header.toString().toLowerCase() === 'authorization') ||
+    (unknownOrigin && header.length === 6 && header.toString().toLowerCase() === 'cookie')
+  )
+}
+
+// https://tools.ietf.org/html/rfc7231#section-6.4
+function cleanRequestHeaders (headers, removeContent, unknownOrigin) {
+  const ret = []
+  if (Array.isArray(headers)) {
+    for (let i = 0; i < headers.length; i += 2) {
+      if (!shouldRemoveHeader(headers[i], removeContent, unknownOrigin)) {
+        ret.push(headers[i], headers[i + 1])
+      }
+    }
+  } else if (headers && typeof headers === 'object') {
+    for (const key of Object.keys(headers)) {
+      if (!shouldRemoveHeader(key, removeContent, unknownOrigin)) {
+        ret.push(key, headers[key])
+      }
+    }
+  } else {
+    assert(headers == null, 'headers must be an object or an array')
+  }
+  return ret
+}
+
+module.exports = RedirectHandler
diff --git a/lib/handler/RetryHandler.js b/lib/handler/RetryHandler.js
new file mode 100644
index 0000000..3710447
--- /dev/null
+++ b/lib/handler/RetryHandler.js
@@ -0,0 +1,336 @@
+const assert = require('assert')
+
+const { kRetryHandlerDefaultRetry } = require('../core/symbols')
+const { RequestRetryError } = require('../core/errors')
+const { isDisturbed, parseHeaders, parseRangeHeader } = require('../core/util')
+
+function calculateRetryAfterHeader (retryAfter) {
+  const current = Date.now()
+  const diff = new Date(retryAfter).getTime() - current
+
+  return diff
+}
+
+class RetryHandler {
+  constructor (opts, handlers) {
+    const { retryOptions, ...dispatchOpts } = opts
+    const {
+      // Retry scoped
+      retry: retryFn,
+      maxRetries,
+      maxTimeout,
+      minTimeout,
+      timeoutFactor,
+      // Response scoped
+      methods,
+      errorCodes,
+      retryAfter,
+      statusCodes
+    } = retryOptions ?? {}
+
+    this.dispatch = handlers.dispatch
+    this.handler = handlers.handler
+    this.opts = dispatchOpts
+    this.abort = null
+    this.aborted = false
+    this.retryOpts = {
+      retry: retryFn ?? RetryHandler[kRetryHandlerDefaultRetry],
+      retryAfter: retryAfter ?? true,
+      maxTimeout: maxTimeout ?? 30 * 1000, // 30s,
+      timeout: minTimeout ?? 500, // .5s
+      timeoutFactor: timeoutFactor ?? 2,
+      maxRetries: maxRetries ?? 5,
+      // What errors we should retry
+      methods: methods ?? ['GET', 'HEAD', 'OPTIONS', 'PUT', 'DELETE', 'TRACE'],
+      // Indicates which errors to retry
+      statusCodes: statusCodes ?? [500, 502, 503, 504, 429],
+      // List of errors to retry
+      errorCodes: errorCodes ?? [
+        'ECONNRESET',
+        'ECONNREFUSED',
+        'ENOTFOUND',
+        'ENETDOWN',
+        'ENETUNREACH',
+        'EHOSTDOWN',
+        'EHOSTUNREACH',
+        'EPIPE'
+      ]
+    }
+
+    this.retryCount = 0
+    this.start = 0
+    this.end = null
+    this.etag = null
+    this.resume = null
+
+    // Handle possible onConnect duplication
+    this.handler.onConnect(reason => {
+      this.aborted = true
+      if (this.abort) {
+        this.abort(reason)
+      } else {
+        this.reason = reason
+      }
+    })
+  }
+
+  onRequestSent () {
+    if (this.handler.onRequestSent) {
+      this.handler.onRequestSent()
+    }
+  }
+
+  onUpgrade (statusCode, headers, socket) {
+    if (this.handler.onUpgrade) {
+      this.handler.onUpgrade(statusCode, headers, socket)
+    }
+  }
+
+  onConnect (abort) {
+    if (this.aborted) {
+      abort(this.reason)
+    } else {
+      this.abort = abort
+    }
+  }
+
+  onBodySent (chunk) {
+    if (this.handler.onBodySent) return this.handler.onBodySent(chunk)
+  }
+
+  static [kRetryHandlerDefaultRetry] (err, { state, opts }, cb) {
+    const { statusCode, code, headers } = err
+    const { method, retryOptions } = opts
+    const {
+      maxRetries,
+      timeout,
+      maxTimeout,
+      timeoutFactor,
+      statusCodes,
+      errorCodes,
+      methods
+    } = retryOptions
+    let { counter, currentTimeout } = state
+
+    currentTimeout =
+      currentTimeout != null && currentTimeout > 0 ? currentTimeout : timeout
+
+    // Any code that is not a Undici's originated and allowed to retry
+    if (
+      code &&
+      code !== 'UND_ERR_REQ_RETRY' &&
+      code !== 'UND_ERR_SOCKET' &&
+      !errorCodes.includes(code)
+    ) {
+      cb(err)
+      return
+    }
+
+    // If a set of method are provided and the current method is not in the list
+    if (Array.isArray(methods) && !methods.includes(method)) {
+      cb(err)
+      return
+    }
+
+    // If a set of status code are provided and the current status code is not in the list
+    if (
+      statusCode != null &&
+      Array.isArray(statusCodes) &&
+      !statusCodes.includes(statusCode)
+    ) {
+      cb(err)
+      return
+    }
+
+    // If we reached the max number of retries
+    if (counter > maxRetries) {
+      cb(err)
+      return
+    }
+
+    let retryAfterHeader = headers != null && headers['retry-after']
+    if (retryAfterHeader) {
+      retryAfterHeader = Number(retryAfterHeader)
+      retryAfterHeader = isNaN(retryAfterHeader)
+        ? calculateRetryAfterHeader(retryAfterHeader)
+        : retryAfterHeader * 1e3 // Retry-After is in seconds
+    }
+
+    const retryTimeout =
+      retryAfterHeader > 0
+        ? Math.min(retryAfterHeader, maxTimeout)
+        : Math.min(currentTimeout * timeoutFactor ** counter, maxTimeout)
+
+    state.currentTimeout = retryTimeout
+
+    setTimeout(() => cb(null), retryTimeout)
+  }
+
+  onHeaders (statusCode, rawHeaders, resume, statusMessage) {
+    const headers = parseHeaders(rawHeaders)
+
+    this.retryCount += 1
+
+    if (statusCode >= 300) {
+      this.abort(
+        new RequestRetryError('Request failed', statusCode, {
+          headers,
+          count: this.retryCount
+        })
+      )
+      return false
+    }
+
+    // Checkpoint for resume from where we left it
+    if (this.resume != null) {
+      this.resume = null
+
+      if (statusCode !== 206) {
+        return true
+      }
+
+      const contentRange = parseRangeHeader(headers['content-range'])
+      // If no content range
+      if (!contentRange) {
+        this.abort(
+          new RequestRetryError('Content-Range mismatch', statusCode, {
+            headers,
+            count: this.retryCount
+          })
+        )
+        return false
+      }
+
+      // Let's start with a weak etag check
+      if (this.etag != null && this.etag !== headers.etag) {
+        this.abort(
+          new RequestRetryError('ETag mismatch', statusCode, {
+            headers,
+            count: this.retryCount
+          })
+        )
+        return false
+      }
+
+      const { start, size, end = size } = contentRange
+
+      assert(this.start === start, 'content-range mismatch')
+      assert(this.end == null || this.end === end, 'content-range mismatch')
+
+      this.resume = resume
+      return true
+    }
+
+    if (this.end == null) {
+      if (statusCode === 206) {
+        // First time we receive 206
+        const range = parseRangeHeader(headers['content-range'])
+
+        if (range == null) {
+          return this.handler.onHeaders(
+            statusCode,
+            rawHeaders,
+            resume,
+            statusMessage
+          )
+        }
+
+        const { start, size, end = size } = range
+
+        assert(
+          start != null && Number.isFinite(start) && this.start !== start,
+          'content-range mismatch'
+        )
+        assert(Number.isFinite(start))
+        assert(
+          end != null && Number.isFinite(end) && this.end !== end,
+          'invalid content-length'
+        )
+
+        this.start = start
+        this.end = end
+      }
+
+      // We make our best to checkpoint the body for further range headers
+      if (this.end == null) {
+        const contentLength = headers['content-length']
+        this.end = contentLength != null ? Number(contentLength) : null
+      }
+
+      assert(Number.isFinite(this.start))
+      assert(
+        this.end == null || Number.isFinite(this.end),
+        'invalid content-length'
+      )
+
+      this.resume = resume
+      this.etag = headers.etag != null ? headers.etag : null
+
+      return this.handler.onHeaders(
+        statusCode,
+        rawHeaders,
+        resume,
+        statusMessage
+      )
+    }
+
+    const err = new RequestRetryError('Request failed', statusCode, {
+      headers,
+      count: this.retryCount
+    })
+
+    this.abort(err)
+
+    return false
+  }
+
+  onData (chunk) {
+    this.start += chunk.length
+
+    return this.handler.onData(chunk)
+  }
+
+  onComplete (rawTrailers) {
+    this.retryCount = 0
+    return this.handler.onComplete(rawTrailers)
+  }
+
+  onError (err) {
+    if (this.aborted || isDisturbed(this.opts.body)) {
+      return this.handler.onError(err)
+    }
+
+    this.retryOpts.retry(
+      err,
+      {
+        state: { counter: this.retryCount++, currentTimeout: this.retryAfter },
+        opts: { retryOptions: this.retryOpts, ...this.opts }
+      },
+      onRetry.bind(this)
+    )
+
+    function onRetry (err) {
+      if (err != null || this.aborted || isDisturbed(this.opts.body)) {
+        return this.handler.onError(err)
+      }
+
+      if (this.start !== 0) {
+        this.opts = {
+          ...this.opts,
+          headers: {
+            ...this.opts.headers,
+            range: `bytes=${this.start}-${this.end ?? ''}`
+          }
+        }
+      }
+
+      try {
+        this.dispatch(this.opts, this)
+      } catch (err) {
+        this.handler.onError(err)
+      }
+    }
+  }
+}
+
+module.exports = RetryHandler