eioc 0.3.1

//! HTTP long-polling transport tasks for Engine.IO v4.

use crate::ENGINE_IO_VERSION;
use crate::engine::FrameSender;
use crate::error::{PollingError, TransportError};
use crate::packet::{Frame, Handshake, Packet};
use crate::prelude::WebSocketStream;
use crate::websocket::WebSocketConnector;
use base64::prelude::{BASE64_STANDARD, Engine as _};
use bytes::Bytes;
use bytestring::ByteString;
use reqwest::Client;
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use url::Url;

const SEPARATOR: char = '\x1e';

impl Frame {
    fn decode_binary(data: &[u8]) -> Result<Self, PollingError> {
        Ok(Self::Binary(Bytes::from(BASE64_STANDARD.decode(data)?)))
    }

    fn decode_packet(bytes: &ByteString) -> Result<Self, PollingError> {
        Ok(Self::Packet(Packet::decode(bytes)?))
    }

    fn decode(bytes: &ByteString) -> Result<Self, PollingError> {
        let mut chars = bytes.chars();

        if chars.next().is_some_and(|b| b == 'b') {
            Self::decode_binary(chars.as_str().as_bytes())
        } else {
            Self::decode_packet(bytes)
        }
    }

    fn write(&self, buffer: &mut String) {
        match self {
            Frame::Packet(packet) => packet.write(buffer),
            Frame::Binary(bytes) => {
                buffer.push('b');
                BASE64_STANDARD.encode_string(bytes, buffer);
            }
        }
    }
}

fn decode_frames(bytes: &ByteString) -> Result<Vec<Frame>, PollingError> {
    bytes
        .split(SEPARATOR)
        .map(|s| Frame::decode(&bytes.slice_ref(s)))
        .collect()
}

fn encode_frames(frames: &[Frame]) -> String {
    let mut buffer = String::new();
    for (i, frame) in frames.iter().enumerate() {
        if i > 0 {
            buffer.push(SEPARATOR);
        }
        frame.write(&mut buffer);
    }
    buffer
}

/// Builds the polling URL by appending the EIO version and transport parameters.
fn polling_url(mut base_url: Url) -> Url {
    base_url
        .query_pairs_mut()
        .append_pair("EIO", ENGINE_IO_VERSION)
        .append_pair("transport", "polling");
    base_url
}

/// Wraps a [`reqwest::Client`] with Engine.IO HTTP polling helpers.
#[derive(Clone)]
pub struct PollingClient(pub Client);

impl PollingClient {
    async fn get(&self, url: &Url) -> Result<Vec<Frame>, PollingError> {
        let response = self
            .0
            .get(url.as_str())
            .send()
            .await?
            .error_for_status()?
            .text()
            .await?;

        tracing::trace!(bytes = response.len(), "<- GET");
        decode_frames(&ByteString::from(response))
    }

    async fn get_one(&self, url: &Url) -> Result<Frame, PollingError> {
        let response = self
            .0
            .get(url.as_str())
            .send()
            .await?
            .error_for_status()?
            .text()
            .await?;

        tracing::trace!(bytes = response.len(), "<- GET");

        Frame::decode(&ByteString::from(response))
    }

    async fn post(&self, url: &Url, frames: &[Frame]) -> Result<(), PollingError> {
        let body = encode_frames(frames);
        tracing::trace!(bytes = body.len(), "-> POST");

        let response = self
            .0
            .post(url.as_str())
            .body(body)
            .send()
            .await?
            .error_for_status()?
            .text()
            .await?;

        if !response.eq_ignore_ascii_case("ok") {
            return Err(PollingError::Response(response));
        }

        Ok(())
    }

    /// Loops batched POST requests until `token` fires or `rx` closes.
    ///
    /// Returns `rx` on exit so the WebSocket phase can reuse it.
    #[tracing::instrument(level = "debug", skip_all, err)]
    async fn post_until_cancelled(
        &self,
        url: &Url,
        mut rx: mpsc::Receiver<Frame>,
        token: CancellationToken,
    ) -> Result<mpsc::Receiver<Frame>, TransportError> {
        let mut buffer = Vec::with_capacity(8);

        loop {
            let count = tokio::select! {
                () = token.cancelled() => {
                    tracing::debug!("cancelled polling POST");
                    break;
                },
                count = rx.recv_many(&mut buffer, 8) => count,
            };

            if count == 0 {
                break;
            }

            self.post(url, &buffer).await?;
            buffer.clear();
        }

        Ok(rx)
    }

    /// Drains all outbound frames from `rx` until the sender is dropped.
    #[tracing::instrument(level = "debug", skip_all, err)]
    async fn post_until_closed(
        &self,
        url: &Url,
        mut rx: mpsc::Receiver<Frame>,
    ) -> Result<(), TransportError> {
        let mut buffer = Vec::with_capacity(8);

        while rx.recv_many(&mut buffer, 8).await > 0 {
            self.post(url, &buffer).await?;
            buffer.clear();
        }

        Ok(())
    }

    /// Loops GET requests, decoding each response and forwarding frames to the engine.
    ///
    /// Exits when `token` fires, the engine channel closes, or an HTTP error occurs.
    #[tracing::instrument(level = "debug", skip_all, err)]
    async fn get_until_cancelled(
        &self,
        url: &Url,
        frame_tx: FrameSender,
        token: CancellationToken,
    ) -> Result<FrameSender, TransportError> {
        while !token.is_cancelled() {
            for frame in self.get(url).await? {
                frame_tx.send(frame).await?;
            }
        }

        tracing::debug!("cancelled polling GET");

        Ok(frame_tx)
    }

    /// Runs the full polling transport lifecycle: handshake, GET/POST loops, and optional WebSocket upgrade.
    ///
    /// # Errors
    ///
    /// Returns an error if a network, protocol, or channel failure occurs.
    #[tracing::instrument(skip_all, err)]
    pub async fn transport<C>(
        self,
        base_url: Url,
        connector: C,
        handshake_tx: oneshot::Sender<Handshake>,
        frame_tx: FrameSender,
        transport_rx: mpsc::Receiver<Frame>,
        token: CancellationToken,
    ) -> Result<(), TransportError>
    where
        C: WebSocketConnector,
    {
        let mut url = polling_url(base_url.clone());

        let span = tracing::debug_span!("connect", %url);

        let handshake = match self.get_one(&url).instrument(span).await? {
            Frame::Packet(Packet::Open(handshake)) => handshake,
            frame => return Err(TransportError::Open(frame)),
        };

        url.query_pairs_mut().append_pair("sid", &handshake.sid);
        let do_upgrade = handshake.can_upgrade_to_websocket();
        let sid = handshake.sid.clone();

        handshake_tx
            .send(handshake)
            .map_err(TransportError::SendHandshake)?;

        let _guard = token.clone().drop_guard();

        let child_token = token.child_token();

        let get_fut = self.get_until_cancelled(&url, frame_tx, child_token.clone());

        if do_upgrade {
            let post_fut = self.post_until_cancelled(&url, transport_rx, child_token.clone());

            let stream_fut = async {
                let _guard = child_token.drop_guard();

                let stream = WebSocketStream::connect(base_url, Some(&sid), connector).await?;

                tracing::debug!("paused polling transport");

                Ok::<_, TransportError>(stream)
            };

            let (get_result, post_result, stream_result) =
                tokio::join!(get_fut, post_fut, stream_fut);

            let frame_tx = get_result?;
            let transport_rx = post_result?;
            let stream = stream_result?;

            stream.transport(None, frame_tx, transport_rx).await?;
        } else {
            let post_fut = self.post_until_closed(&url, transport_rx);

            let (get_result, post_result) = tokio::join!(get_fut, post_fut);

            get_result?;
            post_result?;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn bss(s: &'static str) -> ByteString {
        ByteString::from_static(s)
    }

    #[test]
    fn polling_url_appends_params() {
        let base = Url::parse("http://localhost:3000/socket.io/").unwrap();
        let url = polling_url(base);
        let query = url.query().unwrap();
        assert!(query.contains("EIO=4"));
        assert!(query.contains("transport=polling"));
    }

    #[test]
    fn frame_decode_text_packet() {
        let frame = Frame::decode(&bss("4hello")).unwrap();
        assert!(matches!(frame, Frame::Packet(Packet::Message(m)) if m == "hello"));
    }

    #[test]
    fn frame_decode_binary_base64() {
        use base64::prelude::{BASE64_STANDARD, Engine as _};
        let encoded = BASE64_STANDARD.encode(b"abc");
        let input = ByteString::from(format!("b{encoded}"));
        let frame = Frame::decode(&input).unwrap();
        assert!(matches!(frame, Frame::Binary(b) if b.as_ref() == b"abc"));
    }

    #[test]
    fn frame_decode_invalid_base64_is_error() {
        assert!(Frame::decode(&bss("b!!!")).is_err());
    }

    #[test]
    fn frame_write_packet() {
        let frame = Frame::Packet(Packet::Message("hello".into()));
        let mut buf = String::new();
        frame.write(&mut buf);
        assert_eq!(buf, "4hello");
    }

    #[test]
    fn frame_write_binary() {
        use base64::prelude::{BASE64_STANDARD, Engine as _};
        let raw = Bytes::from_static(b"abc");
        let frame = Frame::Binary(raw);
        let mut buf = String::new();
        frame.write(&mut buf);
        assert_eq!(buf, format!("b{}", BASE64_STANDARD.encode(b"abc")));
    }

    #[test]
    fn decode_encode_frames_roundtrip() {
        let text = bss("4hello\x1e4world");
        let frames = decode_frames(&text).unwrap();
        assert_eq!(frames.len(), 2);
        let encoded = encode_frames(&frames);
        assert_eq!(encoded, "4hello\x1e4world");
    }

    #[test]
    fn encode_frames_single() {
        let frames = vec![Frame::Packet(Packet::Pong("probe".into()))];
        assert_eq!(encode_frames(&frames), "3probe");
    }

    #[test]
    fn encode_frames_empty() {
        assert_eq!(encode_frames(&[]), "");
    }

    #[test]
    fn decode_frames_error_propagates() {
        assert!(decode_frames(&bss("9invalid")).is_err());
    }
}