'use strict'

const { test } = require('node:test')
const { once } = require('node:events')
const { randomFillSync } = require('node:crypto')
const { deflateRawSync } = require('node:zlib')
const { setTimeout: sleep } = require('node:timers/promises')
const { WebSocketServer } = require('ws')
const { WebSocket, Agent } = require('../..')

test('Compressed message under limit decompresses successfully', async (t) => {
  const server = new WebSocketServer({
    port: 0,
    perMessageDeflate: true
  })

  t.after(() => server.close())

  await once(server, 'listening')

  server.on('connection', (ws) => {
    // Send 1 KB of data (well under any reasonable limit)
    ws.send(Buffer.alloc(1024, 0x41), { binary: true })
  })

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

  const [event] = await once(client, 'message')
  t.assert.strictEqual(event.data.size, 1024)
  client.close()
})

test('Agent webSocketOptions.maxPayloadSize is read correctly', async (t) => {
  const customLimit = 128 * 1024 * 1024 // 128 MB
  const agent = new Agent({
    webSocket: {
      maxPayloadSize: customLimit
    }
  })

  t.after(() => agent.close())

  // Verify the option is stored and retrievable
  t.assert.strictEqual(agent.webSocketOptions.maxPayloadSize, customLimit)
})

test('Agent with default webSocketOptions uses 128 MB limit', async (t) => {
  const agent = new Agent()

  t.after(() => agent.close())

  // Default should be 128 MB
  t.assert.strictEqual(agent.webSocketOptions.maxPayloadSize, 128 * 1024 * 1024)
})

test('Custom maxPayloadSize allows messages under limit', async (t) => {
  const server = new WebSocketServer({
    port: 0,
    perMessageDeflate: true
  })

  t.after(() => server.close())
  await once(server, 'listening')

  const dataSize = 512 * 1024 // 512 KB

  server.on('connection', (ws) => {
    ws.send(Buffer.alloc(dataSize, 0x41), { binary: true })
  })

  // Set custom limit of 1 MB via Agent
  const agent = new Agent({
    webSocket: {
      maxPayloadSize: 1 * 1024 * 1024
    }
  })

  t.after(() => agent.close())

  const client = new WebSocket(`ws://127.0.0.1:${server.address().port}`, { dispatcher: agent })

  const [event] = await once(client, 'message')
  t.assert.strictEqual(event.data.size, dataSize, 'Message under limit should be received')
  client.close()
})

test('Messages at exactly the limit succeed', async (t) => {
  const limit = 1 * 1024 * 1024 // 1 MB
  const server = new WebSocketServer({
    port: 0,
    perMessageDeflate: true
  })

  t.after(() => server.close())
  await once(server, 'listening')

  server.on('connection', (ws) => {
    ws.send(Buffer.alloc(limit, 0x41), { binary: true })
  })

  const agent = new Agent({
    webSocket: {
      maxPayloadSize: limit
    }
  })

  t.after(() => agent.close())

  const client = new WebSocket(`ws://127.0.0.1:${server.address().port}`, { dispatcher: agent })

  const [event] = await once(client, 'message')
  t.assert.strictEqual(event.data.size, limit, 'Message at exactly the limit should succeed')
  client.close()
})

test('Compressed frame payload over wire-size limit is rejected', async (t) => {
  const limit = 64 * 1024
  const server = new WebSocketServer({
    port: 0,
    perMessageDeflate: true
  })

  t.after(() => server.close())
  await once(server, 'listening')

  let payload = null
  for (let i = 0; i < 10; i++) {
    const candidate = randomFillSync(Buffer.alloc(limit))
    if (deflateRawSync(candidate).length > limit) {
      payload = candidate
      break
    }
  }

  t.assert.ok(payload, 'Expected incompressible payload with compressed wire size over the limit')

  let messageReceived = false

  server.on('connection', (ws) => {
    ws.send(payload, { binary: true, compress: true })
  })

  const agent = new Agent({
    webSocket: {
      maxPayloadSize: limit
    }
  })

  t.after(() => agent.close())

  const client = new WebSocket(`ws://127.0.0.1:${server.address().port}`, { dispatcher: agent })

  client.addEventListener('message', () => {
    messageReceived = true
  })

  const closePromise = once(client, 'close')
  const timeoutPromise = sleep(5000)

  await Promise.race([closePromise, timeoutPromise])

  t.assert.strictEqual(messageReceived, false, 'Compressed frame over wire-size limit should be rejected')
  t.assert.strictEqual(client.readyState, WebSocket.CLOSED, 'Connection should be closed after exceeding limit')
})

test('Messages over the limit are rejected', async (t) => {
  const limit = 1 * 1024 * 1024 // 1 MB
  const server = new WebSocketServer({
    port: 0,
    perMessageDeflate: true
  })

  t.after(() => server.close())
  await once(server, 'listening')

  let messageReceived = false
  let closeEvent = null

  server.on('connection', (ws) => {
    // Send 2 MB of data, which exceeds the 1 MB limit
    ws.send(Buffer.alloc(2 * 1024 * 1024, 0x41), { binary: true })
  })

  const agent = new Agent({
    webSocket: {
      maxPayloadSize: limit
    }
  })

  t.after(() => agent.close())

  const client = new WebSocket(`ws://127.0.0.1:${server.address().port}`, { dispatcher: agent })

  client.addEventListener('message', () => {
    messageReceived = true
  })

  client.addEventListener('close', (event) => {
    closeEvent = event
  })

  // Wait for connection to close (should happen when limit is exceeded)
  // Use Promise.race with a timeout to avoid hanging forever
  const closePromise = once(client, 'close')
  const timeoutPromise = sleep(5000)

  await Promise.race([closePromise, timeoutPromise])

  t.assert.strictEqual(messageReceived, false, 'Message over limit should be rejected')
  t.assert.ok(closeEvent !== null, 'Close event should have been emitted')
  t.assert.strictEqual(client.readyState, WebSocket.CLOSED, 'Connection should be closed after exceeding limit')
})

test('Limit can be disabled by setting maxPayloadSize to 0', async (t) => {
  const server = new WebSocketServer({
    port: 0,
    perMessageDeflate: true
  })

  t.after(() => server.close())
  await once(server, 'listening')

  // Keep this comfortably above the smaller limits used elsewhere in this file,
  // while avoiding the 100 MB transfer that can be slow on CI.
  const dataSize = 2 * 1024 * 1024 // 2 MB

  server.on('connection', (ws) => {
    ws.send(Buffer.alloc(dataSize, 0x41), { binary: true })
  })

  // Set limit to 0 (disabled)
  const agent = new Agent({
    webSocket: {
      maxPayloadSize: 0
    }
  })

  t.after(() => agent.close())

  const client = new WebSocket(`ws://127.0.0.1:${server.address().port}`, { dispatcher: agent })
  const timeout = Symbol('timeout')

  const result = await Promise.race([
    once(client, 'message'),
    sleep(10000, timeout)
  ])

  t.assert.notStrictEqual(result, timeout, 'Test timed out waiting for large message')
  t.assert.strictEqual(result[0].data.size, dataSize, 'Large message should be received when limit is disabled')
  client.close()
})

test('Fragmented compressed payload over total limit is rejected', async (t) => {
  const limit = 1 * 1024 * 1024 // 1 MB
  const fragmentSize = 768 * 1024 // 768 KB
  const server = new WebSocketServer({
    port: 0,
    perMessageDeflate: true
  })

  t.after(() => server.close())
  await once(server, 'listening')

  let messageReceived = false

  server.on('connection', (ws) => {
    ws.send(Buffer.alloc(fragmentSize, 0x41), {
      binary: true,
      compress: true,
      fin: false
    })

    ws.send(Buffer.alloc(fragmentSize, 0x41), {
      binary: true,
      compress: true,
      fin: true
    })
  })

  const agent = new Agent({
    webSocket: {
      maxPayloadSize: limit
    }
  })

  t.after(() => agent.close())

  const client = new WebSocket(`ws://127.0.0.1:${server.address().port}`, { dispatcher: agent })

  client.addEventListener('message', () => {
    messageReceived = true
  })

  const closePromise = once(client, 'close')
  const timeoutPromise = sleep(5000)

  await Promise.race([closePromise, timeoutPromise])

  t.assert.strictEqual(messageReceived, false, 'Fragmented compressed message over total limit should be rejected')
  t.assert.strictEqual(client.readyState, WebSocket.CLOSED, 'Connection should be closed after exceeding limit')
})

test('Raw uncompressed payload over immediate limit is rejected', async (t) => {
  const limit = 100
  const server = new WebSocketServer({
    port: 0,
    perMessageDeflate: false // Disable compression
  })

  t.after(() => server.close())
  await once(server, 'listening')

  let messageReceived = false

  server.on('connection', (ws) => {
    // Send 101 bytes uncompressed so the inline payload length path is used.
    ws.send(Buffer.alloc(101, 0x41), { binary: true })
  })

  const agent = new Agent({
    webSocket: {
      maxPayloadSize: limit
    }
  })

  t.after(() => agent.close())

  const client = new WebSocket(`ws://127.0.0.1:${server.address().port}`, { dispatcher: agent })

  client.addEventListener('message', () => {
    messageReceived = true
  })

  const closePromise = once(client, 'close')
  const timeoutPromise = sleep(5000)

  await Promise.race([closePromise, timeoutPromise])

  t.assert.strictEqual(messageReceived, false, 'Raw uncompressed message over limit should be rejected')
  t.assert.strictEqual(client.readyState, WebSocket.CLOSED, 'Connection should be closed after exceeding limit')
})

test('Raw uncompressed payload over 16-bit extended limit is rejected', async (t) => {
  const limit = 1 * 1024 // 1 KB
  const server = new WebSocketServer({
    port: 0,
    perMessageDeflate: false // Disable compression
  })

  t.after(() => server.close())
  await once(server, 'listening')

  let messageReceived = false

  server.on('connection', (ws) => {
    // Send 2 KB uncompressed so the extended 16-bit payload length path is used.
    ws.send(Buffer.alloc(2 * 1024, 0x41), { binary: true })
  })

  const agent = new Agent({
    webSocket: {
      maxPayloadSize: limit
    }
  })

  t.after(() => agent.close())

  const client = new WebSocket(`ws://127.0.0.1:${server.address().port}`, { dispatcher: agent })

  client.addEventListener('message', () => {
    messageReceived = true
  })

  const closePromise = once(client, 'close')
  const timeoutPromise = sleep(5000)

  await Promise.race([closePromise, timeoutPromise])

  t.assert.strictEqual(messageReceived, false, 'Raw uncompressed message over limit should be rejected')
  t.assert.strictEqual(client.readyState, WebSocket.CLOSED, 'Connection should be closed after exceeding limit')
})

test('Raw uncompressed payload over 64-bit extended limit is rejected', async (t) => {
  const limit = 1 * 1024 * 1024 // 1 MB
  const server = new WebSocketServer({
    port: 0,
    perMessageDeflate: false // Disable compression
  })

  t.after(() => server.close())
  await once(server, 'listening')

  let messageReceived = false

  server.on('connection', (ws) => {
    // Send 2 MB uncompressed so the extended 64-bit payload length path is used.
    ws.send(Buffer.alloc(2 * 1024 * 1024, 0x41), { binary: true })
  })

  const agent = new Agent({
    webSocket: {
      maxPayloadSize: limit
    }
  })

  t.after(() => agent.close())

  const client = new WebSocket(`ws://127.0.0.1:${server.address().port}`, { dispatcher: agent })

  client.addEventListener('message', () => {
    messageReceived = true
  })

  const closePromise = once(client, 'close')
  const timeoutPromise = sleep(5000)

  await Promise.race([closePromise, timeoutPromise])

  t.assert.strictEqual(messageReceived, false, 'Raw uncompressed message over limit should be rejected')
  t.assert.strictEqual(client.readyState, WebSocket.CLOSED, 'Connection should be closed after exceeding limit')
})

test('cumulative payload size', (t, done) => {
  t.plan(1)

  const LIMIT = 100
  const FRAGMENT_SIZE = 60
  const NUM_FRAGMENTS = 10

  const server = new WebSocketServer({ port: 0 })

  server.on('connection', (ws) => {
    const socket = ws._socket
    const payload = Buffer.alloc(FRAGMENT_SIZE, 0x41)

    for (let i = 0; i < NUM_FRAGMENTS; i++) {
      const fin = i === NUM_FRAGMENTS - 1 ? 0x80 : 0x00
      const opcode = i === 0 ? 0x02 : 0x00
      const header = Buffer.alloc(2)
      header[0] = fin | opcode
      header[1] = FRAGMENT_SIZE
      socket.write(header)
      socket.write(payload)
    }
  })

  const agent = new Agent({
    webSocket: {
      maxPayloadSize: LIMIT
    }
  })

  const client = new WebSocket(`ws://127.0.0.1:${server.address().port}`, { dispatcher: agent })

  t.after(async () => {
    client.close()
    server.close()
    await agent.close()
  })

  client.onmessage = t.assert.fail

  client.addEventListener('error', (event) => {
    t.assert.ok(event)
    done()
  })
})
