1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

//! Traits and types used by retry mechanism.

use std::future::Future;
use std::io;
use std::io::ErrorKind;
use std::marker::PhantomData;
use std::pin::Pin;

use futures::SinkExt;
use futures::{Sink, Stream};
use stream_reconnect::UnderlyingStream;

pub use config::RetryConfig;
pub use context::RetryContext;
pub use policy::BEBIterator;

use crate::errors::StreamError;
use crate::packet::Packet;

mod config;
mod context;
mod policy;

/// Trait of helper objects to connect bilibili websocket server.
///
/// This trait is used when constructing normal bililive streams or auto-retry bililive streams.
///
/// An implementation of `WsStreamTrait` takes in a ws server url and decodes the data into a stream
/// of [`Packet`](crate::packet::Packet) with heartbeat auto-response mechanism implemented
/// (see [`HeartbeatStream`](crate::stream::HeartbeatStream) for details).
#[cfg(feature = "not-send")]
pub trait WsStreamTrait<E> {
    /// The returned stream type.
    type Stream: Stream<Item = Result<Packet, StreamError<E>>>
        + Sink<Packet, Error = StreamError<E>>
        + Unpin
        + Sized;
    /// Connect to bilibili websocket server.
    ///
    /// # Errors
    /// Returns an error when websocket connection fails.
    fn connect(url: &str) -> Pin<Box<dyn Future<Output = Result<Self::Stream, E>> + '_>>;
}

#[cfg(not(feature = "not-send"))]
pub trait WsStreamTrait<E> {
    /// The returned stream type.
    type Stream: Stream<Item = Result<Packet, StreamError<E>>>
        + Sink<Packet, Error = StreamError<E>>
        + Unpin
        + Sized
        + Send;
    /// Connect to bilibili websocket server.
    ///
    /// # Errors
    /// Returns an error when websocket connection fails.
    fn connect(url: &str) -> Pin<Box<dyn Future<Output = Result<Self::Stream, E>> + Send + '_>>;
}

/// Wrapper for types implementing `WsStreamTrait`.
///
/// This type is used to avoid the orphan rule. Exposed for stream type construction.
#[derive(Debug, Default)]
pub struct WsStream<T: WsStreamTrait<E>, E>(PhantomData<(T, E)>);

impl<T, E> WsStream<T, E>
where
    T: WsStreamTrait<E>,
{
    /// Connect to bilibili websocket server.
    ///
    /// # Errors
    /// Returns an error when websocket connection fails.
    pub async fn connect(url: &str) -> Result<T::Stream, E> {
        T::connect(url).await
    }
}

#[allow(clippy::type_complexity)]
impl<T, E> UnderlyingStream<RetryContext, Result<Packet, StreamError<E>>, StreamError<E>>
    for WsStream<T, E>
where
    T: WsStreamTrait<E>,
    E: std::error::Error,
{
    type Stream = T::Stream;

    #[cfg(feature = "not-send")]
    fn establish(
        mut ctor_arg: RetryContext,
    ) -> Pin<Box<dyn Future<Output = Result<Self::Stream, StreamError<E>>>>> {
        Box::pin(async move {
            let server = ctor_arg.get();
            let mut ws = Self::connect(server)
                .await
                .map_err(StreamError::from_ws_error)?;
            ws.send(Packet::new_room_enter(ctor_arg.config())).await?;
            Ok(ws)
        })
    }

    #[cfg(not(feature = "not-send"))]
    fn establish(
        mut ctor_arg: RetryContext,
    ) -> Pin<Box<dyn Future<Output = Result<Self::Stream, StreamError<E>>> + Send>> {
        Box::pin(async move {
            let server = ctor_arg.get();
            let mut ws = Self::connect(server)
                .await
                .map_err(StreamError::from_ws_error)?;
            ws.send(Packet::new_room_enter(ctor_arg.config())).await?;
            Ok(ws)
        })
    }

    fn is_write_disconnect_error(err: &StreamError<E>) -> bool {
        matches!(err, StreamError::WebSocket(_) | StreamError::IO(_))
    }

    fn is_read_disconnect_error(item: &Result<Packet, StreamError<E>>) -> bool {
        if let Err(e) = item {
            Self::is_write_disconnect_error(e)
        } else {
            false
        }
    }

    fn exhaust_err() -> StreamError<E> {
        StreamError::IO(io::Error::new(
            ErrorKind::NotConnected,
            "Disconnected. Connection attempts have been exhausted.",
        ))
    }
}