main
use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message;
use super::WebSocket;
use tokio::sync::{mpsc, oneshot};
use super::frames::write_frames;
use crate::Frame;
use crate::engine::{self, Data, Packet as EnginePacket};
use crate::session::Input;
pub(super) async fn run_upgrade(
input: mpsc::Sender<Input>,
mut socket: WebSocket,
out_tx: mpsc::Sender<Vec<Frame>>,
mut outbound: mpsc::Receiver<Vec<Frame>>,
) {
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(10);
let mut out_tx = Some(out_tx);
let mut probed = false;
let mut upgraded = false;
loop {
tokio::select! {
_ = input.closed() => break,
_ = tokio::time::sleep_until(deadline), if !upgraded => break,
incoming = socket.next() => {
let Some(Ok(message)) = incoming else {
break;
};
match message {
Message::Text(text) => {
let Some(packet) = engine::decode_packet(&text) else {
break;
};
match packet {
EnginePacket::Ping(payload) if payload == "probe" && !upgraded => {
if !write_frames(&mut socket, vec![Frame::Text("3probe".to_string())]).await {
break;
}
probed = true;
}
EnginePacket::Upgrade(_) if probed && !upgraded => {
let Some(tx) = out_tx.take() else { break };
let (result, receiver) = oneshot::channel();
if input.try_send(Input::Upgrade { tx, result }).is_err() {
break;
}
if !matches!(tokio::time::timeout_at(deadline, receiver).await, Ok(Ok(true))) {
break;
}
upgraded = true;
}
other if upgraded => {
if input.try_send(Input::Packets(vec![other])).is_err() {
break;
}
}
_ => {}
}
}
Message::Binary(data) if upgraded => {
if input
.try_send(Input::Packets(vec![EnginePacket::Message(Data::Binary(
data.to_vec(),
))]))
.is_err()
{
break;
}
}
Message::Binary(_) => {}
Message::Ping(_) => {
if !matches!(tokio::time::timeout(std::time::Duration::from_secs(5), socket.flush()).await, Ok(Ok(()))) {
break;
}
}
Message::Pong(_) | Message::Frame(_) => {}
Message::Close(_) => break,
}
}
frames = outbound.recv() => {
match frames {
Some(frames) => {
if !write_frames(&mut socket, frames).await {
break;
}
}
None => break,
}
}
}
}
if upgraded {
let _ =
tokio::time::timeout(std::time::Duration::from_secs(5), input.send(Input::Close)).await;
}
let _ = tokio::time::timeout(
std::time::Duration::from_secs(5),
SinkExt::close(&mut socket),
)
.await;
}