1'use strict'
2
3const { InvalidArgumentError, RequestAbortedError, SocketError } = require('../core/errors')
4const { AsyncResource } = require('async_hooks')
5const util = require('../core/util')
6const { addSignal, removeSignal } = require('./abort-signal')
7const assert = require('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.callback) {
38      throw new RequestAbortedError()
39    }
40
41    this.abort = abort
42    this.context = null
43  }
44
45  onHeaders () {
46    throw new SocketError('bad upgrade', null)
47  }
48
49  onUpgrade (statusCode, rawHeaders, socket) {
50    const { callback, opaque, context } = this
51
52    assert.strictEqual(statusCode, 101)
53
54    removeSignal(this)
55
56    this.callback = null
57    const headers = this.responseHeaders === 'raw' ? util.parseRawHeaders(rawHeaders) : util.parseHeaders(rawHeaders)
58    this.runInAsyncScope(callback, null, null, {
59      headers,
60      socket,
61      opaque,
62      context
63    })
64  }
65
66  onError (err) {
67    const { callback, opaque } = this
68
69    removeSignal(this)
70
71    if (callback) {
72      this.callback = null
73      queueMicrotask(() => {
74        this.runInAsyncScope(callback, null, err, { opaque })
75      })
76    }
77  }
78}
79
80function upgrade (opts, callback) {
81  if (callback === undefined) {
82    return new Promise((resolve, reject) => {
83      upgrade.call(this, opts, (err, data) => {
84        return err ? reject(err) : resolve(data)
85      })
86    })
87  }
88
89  try {
90    const upgradeHandler = new UpgradeHandler(opts, callback)
91    this.dispatch({
92      ...opts,
93      method: opts.method || 'GET',
94      upgrade: opts.protocol || 'Websocket'
95    }, upgradeHandler)
96  } catch (err) {
97    if (typeof callback !== 'function') {
98      throw err
99    }
100    const opaque = opts && opts.opaque
101    queueMicrotask(() => callback(err, { opaque }))
102  }
103}
104
105module.exports = upgrade
106