'use strict'

const { test } = require('node:test')
const { once } = require('node:events')
const { createServer } = require('node:http')
const { createSecureServer } = require('node:http2')

const { tspl } = require('@matteo.collina/tspl')
const { WebSocketServer, WebSocket: WSWebsocket } = require('ws')
const { key, cert } = require('@metcoder95/https-pem')
const { WebSocket, Agent, EnvHttpProxyAgent } = require('../..')
const { runtimeFeatures } = require('../../lib/util/runtime-features')
const { uid } = require('../../lib/web/websocket/constants')

const crypto = runtimeFeatures.has('crypto')
  ? require('node:crypto')
  : null

test('WebSocket connecting to server that isn\'t a Websocket server', (t) => {
  return new Promise((resolve, reject) => {
    const server = createServer({ joinDuplicateHeaders: true }, (req, res) => {
      t.assert.strictEqual(req.headers.connection, 'upgrade')
      t.assert.strictEqual(req.headers.upgrade, 'websocket')
      t.assert.ok(req.headers['sec-websocket-key'])
      t.assert.strictEqual(req.headers['sec-websocket-version'], '13')

      res.end()
      server.unref()
    }).listen(0, () => {
      const ws = new WebSocket(`ws://localhost:${server.address().port}`)

      // Server isn't a websocket server
      ws.onmessage = ws.onopen = reject

      ws.addEventListener('error', ({ error }) => {
        t.assert.ok(error)
        server.close()
        resolve()
      })
    })
  })
})

test('Open event is emitted', (t) => {
  return new Promise((resolve, reject) => {
    const server = new WebSocketServer({ port: 0 })

    server.on('connection', (ws) => {
      ws.close(1000)
    })

    const ws = new WebSocket(`ws://localhost:${server.address().port}`)

    ws.onmessage = ws.onerror = reject
    ws.addEventListener('open', (t) => {
      server.close()
      resolve()
    })
  })
})

test('WebSocket on H2', { skip: crypto == null }, async (t) => {
  const planner = tspl(t, { plan: 6 })
  const server = createSecureServer({ cert, key, allowHTTP1: true, settings: { enableConnectProtocol: true } })
  const wsServer = new WebSocketServer({ noServer: true })
  server.on('stream', (stream, headers) => {
    planner.equal(headers[':method'], 'CONNECT')
    planner.equal(headers[':protocol'], 'websocket')
    planner.equal(headers[':path'], '/')
    planner.equal(headers[':scheme'], 'https')

    stream.respond({
      ':status': 200,
      'sec-websocket-protocol': headers['sec-websocket-protocol'],
      'sec-websocket-accept': crypto.hash('sha1', `${headers['sec-websocket-key']}${uid}`, 'base64')
    })
    const ws = new WSWebsocket(null, null, { autoPong: true })
    ws.setSocket(stream, Buffer.alloc(0), {
      maxPayload: 104857600,
      skipUTF8Validation: false
    })

    wsServer.emit('connection', ws, stream)
  })

  wsServer.on('connection', (ws) => {
    ws.send('hello')
  })

  server.listen(0)
  await once(server, 'listening')

  const dispatcher = new Agent({
    allowH2: true,
    connect: {
      rejectUnauthorized: false
    }
  })
  const ws = new WebSocket(`wss://localhost:${server.address().port}`, { dispatcher, protocols: ['chat'] })

  t.after(() => {
    // Cleanup - Seems that due to the nature of H2 we need to remove the error listener
    // TODO: investigate if this is a bug
    ws.onerror = null
    return new Promise((resolve) => {
      ws.close()
      server.close()
      wsServer.close(() => {
        dispatcher.close().then(resolve)
      })
    })
  })

  ws.onmessage = (evt) => planner.equal(evt.data, 'hello')
  ws.onerror = (err) => {
    planner.fail(err)
  }
  ws.addEventListener('open', () => planner.ok(true))

  await planner.completed
})

test('WebSocket connecting to server that isn\'t a Websocket server (h2 - supports extended CONNECT protocol)', async (t) => {
  const planner = tspl(t, { plan: 6 })
  const sessions = new Set()
  const h2Server = createSecureServer({ cert, key, settings: { enableConnectProtocol: true } })
    .on('session', (session) => {
      sessions.add(session)
      session.on('close', () => sessions.delete(session))
    })
    .on('stream', (stream, headers) => {
      planner.equal(headers[':method'], 'CONNECT')
      planner.equal(headers[':protocol'], 'websocket')
      planner.ok(headers['sec-websocket-key'])
      planner.equal(headers['sec-websocket-protocol'], 'chat')
      planner.equal(headers['sec-websocket-version'], '13')

      stream.respond({ ':status': 200 })
      stream.close(8) // NGHTTP2_CANCEL
    })
    .listen(0)

  await once(h2Server, 'listening')

  const dispatcher = new Agent({
    allowH2: true,
    connect: {
      rejectUnauthorized: false
    }
  })
  const ws = new WebSocket(`wss://localhost:${h2Server.address().port}`, { dispatcher, protocols: ['chat'] })
  const cleaner = setupListener()

  t.after(async () => {
    cleaner()
    ws.onerror = null
    ws.close()

    for (const session of sessions) {
      session.close()
    }

    await new Promise((resolve) => h2Server.close(resolve))
    await dispatcher.close()
  })

  ws.onmessage = ws.onopen = () => planner.fail('should not open')

  await planner.completed

  function setupListener () {
    ws.addEventListener('error', listener)

    return () => { ws.removeEventListener('error', listener) }

    function listener ({ error }) {
      planner.ok(error)
    }
  }
})

// A regression here hangs rather than fails, so bound the test.
test('WebSocket over H2 fails the handshake when the extended CONNECT response is not 200', { timeout: 10000 }, async (t) => {
  const planner = tspl(t, { plan: 3 })
  const sessions = new Set()
  const h2Server = createSecureServer({ cert, key, settings: { enableConnectProtocol: true } })
    .on('session', (session) => {
      sessions.add(session)
      session.on('close', () => sessions.delete(session))
    })
    .on('stream', (stream) => {
      stream.respond({ ':status': 401 })
      stream.end()
    })
    .listen(0)

  await once(h2Server, 'listening')

  const dispatcher = new Agent({
    allowH2: true,
    connect: {
      rejectUnauthorized: false
    }
  })
  const ws = new WebSocket(`wss://localhost:${h2Server.address().port}`, { dispatcher })

  t.after(async () => {
    for (const session of sessions) {
      session.close()
    }

    await new Promise((resolve) => h2Server.close(resolve))
    await dispatcher.close()
  })

  ws.onmessage = ws.onopen = () => planner.fail('should not open')

  ws.addEventListener('error', ({ error }) => {
    planner.ok(error)
  })

  ws.addEventListener('close', ({ code, wasClean }) => {
    planner.equal(code, 1006)
    planner.equal(wasClean, false)
  })

  await planner.completed
})

test('WebSocket on H2 with a server that does not support extended CONNECT protocol', async (t) => {
  const planner = tspl(t, { plan: 1 })
  const h2Server = createSecureServer({ cert, key, settings: { enableConnectProtocol: false } })
    .on('stream', (stream) => {
      stream.respond({ ':status': 200 })
      stream.end('')
      h2Server.unref()
    })
    .listen(0)

  await once(h2Server, 'listening')
  t.after(() => { h2Server.close() })

  const dispatcher = new Agent({
    allowH2: true,
    connect: {
      rejectUnauthorized: false
    }
  })
  const ws = new WebSocket(`wss://localhost:${h2Server.address().port}`, { dispatcher, protocols: ['chat'] })

  t.after(() => { return ws.close() || dispatcher.close() })

  ws.onmessage = ws.onopen = () => planner.fail('should not open')
  ws.addEventListener('error', ({ error }) => {
    planner.ok(error)
    ws.close()
    h2Server.close()
  })

  await planner.completed
})

test('WebSocket falls back to HTTP/1.1 upgrade when H2 extended CONNECT is unavailable', async (t) => {
  const planner = tspl(t, { plan: 4 })
  const server = createSecureServer({ cert, key, allowHTTP1: true, settings: { enableConnectProtocol: false } })
  const wsServer = new WebSocketServer({ server })

  server.on('stream', () => planner.fail('should not attempt websocket over HTTP/2'))
  server.on('upgrade', (req) => {
    planner.equal(req.httpVersion, '1.1')
  })

  wsServer.on('connection', (ws, req) => {
    planner.equal(req.httpVersion, '1.1')
    ws.send('hello')
  })

  server.listen(0)
  await once(server, 'listening')

  const dispatcher = new EnvHttpProxyAgent({
    noProxy: '*',
    connect: {
      rejectUnauthorized: false
    }
  })
  const ws = new WebSocket(`wss://localhost:${server.address().port}`, { dispatcher })

  t.after(async () => {
    ws.onerror = null
    ws.close()
    await dispatcher.close()
    await new Promise((resolve) => server.close(resolve))
    await new Promise((resolve) => wsServer.close(resolve))
  })

  ws.onerror = ({ error }) => planner.fail(error)
  ws.addEventListener('open', () => {
    planner.ok(true)
  })
  ws.addEventListener('message', (evt) => {
    planner.equal(evt.data, 'hello')
    ws.close()
  })

  await planner.completed
})

test('WebSocket falls back to HTTP/1.1 upgrade when H2 CONNECT support is omitted', async (t) => {
  const planner = tspl(t, { plan: 4 })
  const server = createSecureServer({ cert, key, allowHTTP1: true })
  const wsServer = new WebSocketServer({ server })

  server.on('stream', () => planner.fail('should not attempt websocket over HTTP/2'))
  server.on('upgrade', (req) => {
    planner.equal(req.httpVersion, '1.1')
  })

  wsServer.on('connection', (ws, req) => {
    planner.equal(req.httpVersion, '1.1')
    ws.send('hello')
  })

  server.listen(0)
  await once(server, 'listening')

  const dispatcher = new EnvHttpProxyAgent({
    noProxy: '*',
    connect: {
      rejectUnauthorized: false
    }
  })
  const ws = new WebSocket(`wss://localhost:${server.address().port}`, { dispatcher })

  t.after(async () => {
    ws.onerror = null
    ws.close()
    await dispatcher.close()
    await new Promise((resolve) => server.close(resolve))
    await new Promise((resolve) => wsServer.close(resolve))
  })

  ws.onerror = ({ error }) => planner.fail(error)
  ws.addEventListener('open', () => {
    planner.ok(true)
  })
  ws.addEventListener('message', (evt) => {
    planner.equal(evt.data, 'hello')
    ws.close()
  })

  await planner.completed
})

test('Multiple protocols are joined by a comma', (t) => {
  return new Promise((resolve, reject) => {
    const server = new WebSocketServer({ port: 0 })

    server.on('connection', (ws, req) => {
      t.assert.strictEqual(req.headers['sec-websocket-protocol'], 'chat, echo')

      ws.close(1000)
      server.close()
      resolve()
    })

    const ws = new WebSocket(`ws://localhost:${server.address().port}`, ['chat', 'echo'])
    ws.addEventListener('open', () => ws.close())
  })
})

test('Server doesn\'t send Sec-WebSocket-Protocol header when protocols are used', (t) => {
  return new Promise((resolve, reject) => {
    const server = createServer({ joinDuplicateHeaders: true }, (req, res) => {
      res.statusCode = 101

      req.socket.destroy()
    }).listen(0, () => {
      const ws = new WebSocket(`ws://localhost:${server.address().port}`, 'chat')

      ws.onopen = reject

      ws.addEventListener('error', ({ error }) => {
        t.assert.ok(error)
        server.close()
        resolve()
      })
    })
  })
})

test('Server sends invalid Upgrade header', (t) => {
  return new Promise((resolve, reject) => {
    const server = createServer({ joinDuplicateHeaders: true }, (req, res) => {
      res.setHeader('Upgrade', 'NotWebSocket')
      res.statusCode = 101

      req.socket.destroy()
    }).listen(0, () => {
      const ws = new WebSocket(`ws://localhost:${server.address().port}`)

      ws.onopen = reject

      ws.addEventListener('error', ({ error }) => {
        t.assert.ok(error)
        server.close()
        resolve()
      })
    })
  })
})

test('Server sends invalid Connection header', (t) => {
  return new Promise((resolve, reject) => {
    const server = createServer({ joinDuplicateHeaders: true }, (req, res) => {
      res.setHeader('Upgrade', 'websocket')
      res.setHeader('Connection', 'downgrade')
      res.statusCode = 101

      req.socket.destroy()
    }).listen(0, () => {
      const ws = new WebSocket(`ws://localhost:${server.address().port}`)

      ws.onopen = reject

      ws.addEventListener('error', ({ error }) => {
        t.assert.ok(error)
        server.close()
        resolve()
      })
    })
  })
})

test('Server sends invalid Sec-WebSocket-Accept header', (t) => {
  return new Promise((resolve, reject) => {
    const server = createServer({ joinDuplicateHeaders: true }, (req, res) => {
      res.setHeader('Upgrade', 'websocket')
      res.setHeader('Connection', 'upgrade')
      res.setHeader('Sec-WebSocket-Accept', 'abc')
      res.statusCode = 101

      req.socket.destroy()
    }).listen(0, () => {
      const ws = new WebSocket(`ws://localhost:${server.address().port}`)

      ws.onopen = reject

      ws.addEventListener('error', ({ error }) => {
        t.assert.ok(error)
        server.close()
        resolve()
      })
    })
  })
})

test('Server sends invalid Sec-WebSocket-Extensions header', { skip: runtimeFeatures.has('crypto') === false }, (t) => {
  return new Promise((resolve, reject) => {
    const uid = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'

    const server = createServer({ joinDuplicateHeaders: true }, (req, res) => {
      const key = req.headers['sec-websocket-key']
      t.assert.ok(key)

      const accept = require('node:crypto').hash('sha1', key + uid, 'base64')

      res.setHeader('Upgrade', 'websocket')
      res.setHeader('Connection', 'upgrade')
      res.setHeader('Sec-WebSocket-Accept', accept)
      res.setHeader('Sec-WebSocket-Extensions', 'InvalidExtension')
      res.statusCode = 101

      res.end()
    }).listen(0, () => {
      const ws = new WebSocket(`ws://localhost:${server.address().port}`)

      ws.onopen = reject

      ws.addEventListener('error', ({ error }) => {
        t.assert.ok(error)
        server.close()
        resolve()
      })
    })
  })
})

test('Server sends invalid Sec-WebSocket-Extensions header', { skip: runtimeFeatures.has('crypto') === false }, (t) => {
  const uid = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'

  return new Promise((resolve, reject) => {
    const server = createServer({ joinDuplicateHeaders: true }, (req, res) => {
      const key = req.headers['sec-websocket-key']
      t.assert.ok(key)

      const accept = require('node:crypto').hash('sha1', key + uid, 'base64')

      res.setHeader('Upgrade', 'websocket')
      res.setHeader('Connection', 'upgrade')
      res.setHeader('Sec-WebSocket-Accept', accept)
      res.setHeader('Sec-WebSocket-Protocol', 'echo') // <--
      res.statusCode = 101

      res.end()
    }).listen(0, () => {
      const ws = new WebSocket(`ws://localhost:${server.address().port}`, 'chat')

      ws.onopen = reject

      ws.addEventListener('error', ({ error }) => {
        t.assert.ok(error)
        server.close()
        resolve()
      })
    })
  })
})

test('Server sends Sec-WebSocket-Protocol header when no protocols were requested', { skip: runtimeFeatures.has('crypto') === false }, (t) => {
  const uid = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'

  return new Promise((resolve, reject) => {
    const server = createServer({ joinDuplicateHeaders: true }, (req, res) => {
      const key = req.headers['sec-websocket-key']
      t.assert.ok(key)
      t.assert.strictEqual(req.headers['sec-websocket-protocol'], undefined)

      const accept = require('node:crypto').hash('sha1', key + uid, 'base64')

      res.setHeader('Upgrade', 'websocket')
      res.setHeader('Connection', 'upgrade')
      res.setHeader('Sec-WebSocket-Accept', accept)
      res.setHeader('Sec-WebSocket-Protocol', 'chat') // <-- never requested
      res.statusCode = 101

      res.end()
    }).listen(0, () => {
      // No subprotocols requested, so the request omits Sec-WebSocket-Protocol.
      const ws = new WebSocket(`ws://localhost:${server.address().port}`)

      ws.onopen = reject

      ws.addEventListener('error', ({ error }) => {
        t.assert.ok(error)
        server.close()
        resolve()
      })
    })
  })
})
