Building a real-time chat using WebSockets over HTTP/2 streams

dpx sv create frontend
cd frontend
deno install
deno run dev

<div class="bg-gray-100 flex flex-col h-screen w-screen">
    <div class="text-center text-gray-500 text-sm py-2">
        <h5 class="font-bold text-blue-900">Real-time chat! Click on connection button and wait for someone else to join.</h5>
    </div>
    <div
        class="bg-white border border-gray-300 flex flex-col grow max-w-4xl mx-auto shadow-lg w-full"
    >
    </div>
</div>

<div class="h-full overflow-y-auto p-4 space-y-2">
  {#each chatHistory as chat}
    {#if chat.type === 'received'}
      {@render receivedMessage(chat.text)}
    {:else}
      {@render sentMessage(chat.text)}
    {/if}
  {/each}
</div>
<div class="bg-gray-50 border-t border-gray-300 p-4">
  <div class="flex flex-row gap-4 hidden items-center">
    {@render connectionButton()}
    {@render messageField()}
    {@render sendButton()}
  </div>
</div>

{#snippet connectionButton()}
    <button
        class="{connectionProps.connectionButtonBg} cursor-pointer px-4 py-2 rounded-lg text-white transition sm:w-auto w-full"
        onclick={connection}>{connectionProps.connectionButtonLabel}</button
    >
{/snippet}

{#snippet messageField()}
    <input
        class="border border-gray-300 disabled:cursor-not-allowed grow p-2 rounded-lg w-full"
        disabled={connectionProps.areElementsDisabled}
        onkeydown={(e) => {
            if (e.key === 'Enter') {
                send();
            }
        }}
        placeholder={connectionProps.messageFieldPlaceholder}
        type="text"
        bind:value={messageFieldValue}
    />
{/snippet}

{#snippet receivedMessage(message: string)}
    <div class="flex justify-start">
        <div class="bg-gray-200 max-w-xs px-4 py-2 rounded-lg text-black">
            {message}
        </div>
    </div>
{/snippet}

{#snippet sendButton()}
    <button
        class="bg-blue-500 enabled:hover:bg-blue-600 disabled:cursor-not-allowed cursor-pointer disabled:opacity-50 px-4 py-2 rounded-lg text-white transition sm:w-auto w-full"
        disabled={connectionProps.areElementsDisabled}
        onclick={send}>Send</button
    >
{/snippet}

{#snippet sentMessage(message: string)}
    <div class="flex justify-end">
        <div class="bg-blue-500 max-w-xs px-4 py-2 rounded-lg text-white">
            {message}
        </div>
    </div>
{/snippet}

<script lang="ts">
    interface ConnectionProps {
        areElementsDisabled: boolean;
        connectionButtonBg: string;
        connectionButtonLabel: string;
        messageFieldPlaceholder: string;
    }

    const connect: ConnectionProps = {
        areElementsDisabled: true,
        connectionButtonBg: 'bg-green-500 hover:bg-green-600',
        connectionButtonLabel: 'Connect',
        messageFieldPlaceholder: "Click on 'Connect' to start a conversation"
    };
    const disconect: ConnectionProps = {
        areElementsDisabled: false,
        connectionButtonBg: 'bg-red-500 hover:bg-red-600',
        connectionButtonLabel: 'Disconnect',
        messageFieldPlaceholder: 'Connected! Type your message'
    };
    const waiting: ConnectionProps = {
        areElementsDisabled: true,
        connectionButtonBg: 'bg-blue-700 hover:bg-blue-800',
        connectionButtonLabel: 'Waiting',
        messageFieldPlaceholder: 'Waiting for someone to show up...'
    };

    let chatHistory: { type: 'received' | 'sent'; text: string }[] = $state([]);
    let connectionProps: ConnectionProps = $state(connect);
    let messageFieldValue = $state('');
    let ws: WebSocket | undefined = undefined;

    const connection = () => {
        if (ws === undefined) {
            ws = new WebSocket('wss://localhost:9000/chat');
            ws.addEventListener('close', () => {
                connectionProps = connect;
                ws = undefined;
            });
            ws.addEventListener('message', (event) => {
                if (event.data == 'OK') {
                    connectionProps = disconect;
                    return;
                }
                chatHistory.push({ type: 'received', text: event.data });
            });
            ws.addEventListener('open', () => {
                connectionProps = waiting;
            });
        } else {
            ws.close();
            connectionProps = connect;
        }
    };

    const send = () => {
        if (ws === undefined || messageFieldValue === '') {
            return;
        }
        ws.send(messageFieldValue);
        chatHistory.push({ type: 'sent', text: messageFieldValue });
        messageFieldValue = '';
    };
</script>

[dependencies]
tokio = { default-features = false, features = ["macros", "rt-multi-thread"], version = "1.0" }
tokio-rustls = { default-features = false, features = ["ring"], version = "0.26" }
wtx = { default-features = false, features = ["http-server-framework", "nightly", "tokio", "tokio-rustls", "web-socket"], version = "0.24" }

[package]
edition = "2021"
name = "backend"
version = "0.1.0"

CERTS_DIR="/some/directory/of/your/choice"
openssl req -newkey rsa:2048 -nodes -subj "/C=FI/CN=vahid" -keyout $CERTS_DIR/key.pem -out $CERTS_DIR/key.csr
openssl x509 -signkey $CERTS_DIR/key.pem -in $CERTS_DIR/key.csr -req -days 365 -out $CERTS_DIR/cert.pem
openssl req -x509 -sha256 -nodes -subj "/C=FI/CN=vahid" -days 365 -newkey rsa:2048 -keyout $CERTS_DIR/root-ca.key -out $CERTS_DIR/root-ca.crt
cat <<'EOF' >> $CERTS_DIR/localhost.ext
authorityKeyIdentifier=keyid,issuer
basicConstraints=CA:FALSE
subjectAltName = @alt_names
[alt_names]
DNS.1 = localhost
IP.1 = 127.0.0.1
EOF
openssl x509 -req -CA $CERTS_DIR/root-ca.crt -CAkey $CERTS_DIR/root-ca.key -in $CERTS_DIR/key.csr -out $CERTS_DIR/cert.pem -days 365 -CAcreateserial -extfile $CERTS_DIR/localhost.ext
rm $CERTS_DIR/key.csr
rm $CERTS_DIR/localhost.ext
rm $CERTS_DIR/root-ca.srl

static CERT: &[u8] = include_bytes!("/some/directory/of/your/choice/cert.pem");
static KEY: &[u8] = include_bytes!("/some/directory/of/your/choice/key.pem");
static USER_POOL: LazyLock<Mutex<UserPool>> = LazyLock::new(|| {
  Mutex::const_new(UserPool { matching: Deque::new(), messages: HashMap::new() })
});

#[tokio::main]
async fn main() -> wtx::Result<()> {
  let router = Router::paths(wtx::paths!(("/chat", web_socket(chat)),))?;
  ServerFrameworkBuilder::new(router)
    .enable_connect_protocol(true)
    .max_hpack_len((128 * 1024, 128 * 1024))
    .without_aux()
    .tokio_rustls(
      (CERT, KEY),
      "0.0.0.0:9000",
      Xorshift64::from(simple_seed()),
      |error| eprintln!("{error:?}"),
      |_| Ok(()),
    )
    .await
}

#[derive(Debug)]
struct UserPool {
  matching: Deque<(u128, Waker)>,
  messages: HashMap<u128, (u128, String, Waker)>,
}

async fn chat(manual_stream: ManualStream<(), ServerStream, ()>) -> wtx::Result<()> {
  let wos = WebSocketOverStream::new(
    &Headers::new(),
    false,
    Xorshift64::from(simple_seed()),
    manual_stream.stream,
  )
  .await?;
  ws::exchange_messages(wos).await?;
  Ok(())
}

async fn handshake(
  local_id: u128,
  wos: &mut WebSocketOverStream<ServerStream>,
) -> wtx::Result<u128> {
  let mut user_pin = pin!(USER_POOL.lock());
  let remote_id = poll_fn(|cx| {
    let mut user_guard = ready!(user_pin.as_mut().poll(cx));
    user_pin.set(USER_POOL.lock());
    if let Some((remote_id, _, _)) = user_guard.messages.get(&local_id) {
      return Poll::Ready(Ok(*remote_id));
    }
    if let Some((remote_id, remote_waker)) = user_guard.matching.pop_front() {
      drop(user_guard.messages.insert(local_id, (remote_id, String::new(), NOOP_WAKER.clone())));
      drop(user_guard.messages.insert(remote_id, (local_id, String::new(), NOOP_WAKER.clone())));
      remote_waker.wake();
      Poll::Ready(wtx::Result::Ok(remote_id))
    } else {
      user_guard.matching.push_back((local_id, cx.waker().clone()))?;
      Poll::Pending
    }
  })
  .await?;
  wos.write_frame(&mut Frame::new_fin(OpCode::Text, *b"OK")).await?;
  Ok(remote_id)
}

async fn connection(
  (local_id, remote_id): (u128, u128),
  wos: &mut WebSocketOverStream<ServerStream>,
) -> wtx::Result<()> {
  let mut buffer = Vector::new();
  loop {
    buffer.clear();
    let mut user_pin = pin!(USER_POOL.lock());
    let message_fut = poll_fn(|cx| {
      let mut user_guard = ready!(user_pin.as_mut().poll(cx));
      user_pin.set(USER_POOL.lock());
      let Some((_, message, waker)) = user_guard.messages.get_mut(&local_id) else {
        return Poll::Ready(Err(wtx::Error::ClosedConnection));
      };
      if message.is_empty() {
        waker.clone_from(cx.waker());
        return Poll::Pending;
      }
      Poll::Ready(wtx::Result::Ok(mem::take(message)))
    });
    tokio::select! {
      frame_rslt = wos.read_frame(&mut buffer) => {
        let frame = frame_rslt?;
        match frame.op_code() {
          OpCode::Text => {
            let Some(text) = frame.text_payload() else {
              return Err(wtx::web_socket::WebSocketError::UnexpectedFrame.into());
            };
            let mut user_guard = USER_POOL.lock().await;
            let Some((_, message, waker)) = user_guard.messages.get_mut(&remote_id) else {
              return Err(wtx::Error::ClosedConnection);
            };
            message.push_str(text);
            waker.wake_by_ref();
          }
          OpCode::Close => break,
          _ => {}
        }
      }
      message_rslt = message_fut => {
        wos.write_frame(&mut Frame::new_fin(OpCode::Text, message_rslt?.into_bytes())).await?;
      }
    }
  }
  Ok(())
}

pub(crate) async fn exchange_messages(
  mut wos: WebSocketOverStream<ServerStream>,
) -> wtx::Result<()> {
  let local_id = GenericTime::timestamp()?.as_nanos();
  let remote_id = handshake(local_id, &mut wos).await?;
  let rslt = connection((local_id, remote_id), &mut wos).await;
  wos.close().await?;
  let mut user_guard = USER_POOL.lock().await;
  drop(user_guard.messages.remove(&local_id));
  if let Some((_, _, waker)) = user_guard.messages.remove(&remote_id) {
    waker.wake();
  }
  drop(user_guard);
  rslt
}