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

use super::connection::WebsocketConnectionBus;
use crate::{
    message::connection::{WebsocketRecv, WebsocketSend},
    service::WebsocketService,
};
use lifeline::{dyn_bus::DynBus, prelude::*};
use log::*;
use serde::{de::DeserializeOwned, Serialize};
use tokio::stream::StreamExt;
use tokio::sync::broadcast;

pub struct WebsocketCarrier {
    _websocket: WebsocketService,
    _websocket_send: Lifeline,
    _websocket_recv: Lifeline,
}

pub trait WebsocketMessageBus: Sized {
    type Send: Message<Self, Channel = broadcast::Sender<Self::Send>>
        + Clone
        + Send
        + Sync
        + Serialize
        + 'static;

    type Recv: Message<Self, Channel = broadcast::Sender<Self::Recv>>
        + Clone
        + DeserializeOwned
        + Send
        + Sync
        + 'static;
}

// TODO: why is dynbus required here?? super confusing
impl<B: DynBus> CarryFrom<B> for WebsocketConnectionBus
where
    B: WebsocketMessageBus,
{
    type Lifeline = anyhow::Result<WebsocketCarrier>;

    fn carry_from(&self, bus: &B) -> Self::Lifeline {
        use tungstenite::Message as TungsteniteMessage;

        let _websocket = WebsocketService::spawn(&self)?;
        self.capacity::<WebsocketSend>(512)?;
        self.capacity::<WebsocketRecv>(512)?;

        let _websocket_send = {
            let mut rx = bus.rx::<B::Send>()?;
            let mut tx = self.tx::<WebsocketSend>()?;

            Self::try_task("forward_send", async move {
                while let Some(msg) = rx.recv().await {
                    trace!("send message: {:?}", &msg);
                    match bincode::serialize(&msg) {
                        Ok(vec) => {
                            let send = tx
                                .send(WebsocketSend(TungsteniteMessage::Binary(vec)))
                                .await;

                            if let Err(_e) = send {
                                debug!("sender disconnected - aborting carry.");
                                break;
                            }
                        }
                        Err(e) => error!("failed to send websocket msg: {}", e),
                    };
                }

                tx.send(WebsocketSend(TungsteniteMessage::Close(None)))
                    .await
                    .ok();

                Ok(())
            })
        };

        let _websocket_recv = {
            let mut rx = self.rx::<WebsocketRecv>()?;
            let mut tx = bus.tx::<B::Recv>()?;

            Self::try_task("forward_recv", async move {
                while let Some(msg) = rx.next().await {
                    let data = msg.0.into_data();
                    match bincode::deserialize(data.as_slice()) {
                        Ok(message) => {
                            trace!("recv message: {:?}", &message);
                            tx.send(message).await?;
                        }
                        Err(e) => error!("failed to recv websocket msg: {}", e),
                    };
                }

                Ok(())
            })
        };

        Ok(WebsocketCarrier {
            _websocket,
            _websocket_send,
            _websocket_recv,
        })
    }
}