api-upgrade.js 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105
  1. 'use strict'
  2. const { InvalidArgumentError, RequestAbortedError, SocketError } = require('../core/errors')
  3. const { AsyncResource } = require('async_hooks')
  4. const util = require('../core/util')
  5. const { addSignal, removeSignal } = require('./abort-signal')
  6. const assert = require('assert')
  7. class UpgradeHandler extends AsyncResource {
  8. constructor (opts, callback) {
  9. if (!opts || typeof opts !== 'object') {
  10. throw new InvalidArgumentError('invalid opts')
  11. }
  12. if (typeof callback !== 'function') {
  13. throw new InvalidArgumentError('invalid callback')
  14. }
  15. const { signal, opaque, responseHeaders } = opts
  16. if (signal && typeof signal.on !== 'function' && typeof signal.addEventListener !== 'function') {
  17. throw new InvalidArgumentError('signal must be an EventEmitter or EventTarget')
  18. }
  19. super('UNDICI_UPGRADE')
  20. this.responseHeaders = responseHeaders || null
  21. this.opaque = opaque || null
  22. this.callback = callback
  23. this.abort = null
  24. this.context = null
  25. addSignal(this, signal)
  26. }
  27. onConnect (abort, context) {
  28. if (!this.callback) {
  29. throw new RequestAbortedError()
  30. }
  31. this.abort = abort
  32. this.context = null
  33. }
  34. onHeaders () {
  35. throw new SocketError('bad upgrade', null)
  36. }
  37. onUpgrade (statusCode, rawHeaders, socket) {
  38. const { callback, opaque, context } = this
  39. assert.strictEqual(statusCode, 101)
  40. removeSignal(this)
  41. this.callback = null
  42. const headers = this.responseHeaders === 'raw' ? util.parseRawHeaders(rawHeaders) : util.parseHeaders(rawHeaders)
  43. this.runInAsyncScope(callback, null, null, {
  44. headers,
  45. socket,
  46. opaque,
  47. context
  48. })
  49. }
  50. onError (err) {
  51. const { callback, opaque } = this
  52. removeSignal(this)
  53. if (callback) {
  54. this.callback = null
  55. queueMicrotask(() => {
  56. this.runInAsyncScope(callback, null, err, { opaque })
  57. })
  58. }
  59. }
  60. }
  61. function upgrade (opts, callback) {
  62. if (callback === undefined) {
  63. return new Promise((resolve, reject) => {
  64. upgrade.call(this, opts, (err, data) => {
  65. return err ? reject(err) : resolve(data)
  66. })
  67. })
  68. }
  69. try {
  70. const upgradeHandler = new UpgradeHandler(opts, callback)
  71. this.dispatch({
  72. ...opts,
  73. method: opts.method || 'GET',
  74. upgrade: opts.protocol || 'Websocket'
  75. }, upgradeHandler)
  76. } catch (err) {
  77. if (typeof callback !== 'function') {
  78. throw err
  79. }
  80. const opaque = opts && opts.opaque
  81. queueMicrotask(() => callback(err, { opaque }))
  82. }
  83. }
  84. module.exports = upgrade