codekingpro/portable-devtools

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