playwright/imp/
websocket.rs

1use crate::imp::{core::*, prelude::*};
2
3#[derive(Debug)]
4pub(crate) struct WebSocket {
5    channel: ChannelOwner,
6    url: String,
7    var: Mutex<Variable>,
8    tx: Mutex<Option<broadcast::Sender<Evt>>>
9}
10
11#[derive(Debug, Default)]
12struct Variable {
13    is_closed: bool
14}
15
16impl WebSocket {
17    pub(crate) fn try_new(channel: ChannelOwner) -> Result<Self, Error> {
18        let Initializer { url } = serde_json::from_value(channel.initializer.clone())?;
19        Ok(Self {
20            channel,
21            url,
22            var: Mutex::default(),
23            tx: Mutex::default()
24        })
25    }
26
27    pub(crate) fn url(&self) -> &str { &self.url }
28}
29
30impl WebSocket {
31    pub(crate) fn is_closed(&self) -> bool { self.var.lock().unwrap().is_closed }
32
33    fn on_frame_sent(&self, params: Map<String, Value>) -> Result<(), Error> {
34        let buffer = parse_frame(params)?;
35        self.emit_event(Evt::FrameSent(buffer));
36        Ok(())
37    }
38
39    fn on_frame_received(&self, params: Map<String, Value>) -> Result<(), Error> {
40        let buffer = parse_frame(params)?;
41        self.emit_event(Evt::FrameReceived(buffer));
42        Ok(())
43    }
44}
45
46fn parse_frame(params: Map<String, Value>) -> Result<Buffer, Error> {
47    #[derive(Deserialize)]
48    struct De {
49        opcode: i32,
50        data: String
51    }
52    let De { opcode, data } = serde_json::from_value(params.into())?;
53    let buffer = if opcode == 2 {
54        let bytes = base64::decode(data).map_err(Error::InvalidBase64)?;
55        Buffer::Bytes(bytes)
56    } else {
57        Buffer::String(data)
58    };
59    Ok(buffer)
60}
61
62impl RemoteObject for WebSocket {
63    fn channel(&self) -> &ChannelOwner { &self.channel }
64    fn channel_mut(&mut self) -> &mut ChannelOwner { &mut self.channel }
65
66    fn handle_event(
67        &self,
68        _ctx: &Context,
69        method: Str<Method>,
70        params: Map<String, Value>
71    ) -> Result<(), Error> {
72        match method.as_str() {
73            "framesent" => self.on_frame_sent(params)?,
74            "framereceived" => self.on_frame_received(params)?,
75            "error" => {
76                let error: Value = params.get("error").cloned().unwrap_or_default();
77                self.emit_event(Evt::Error(error));
78            }
79            "close" => {
80                self.var.lock().unwrap().is_closed = true;
81                self.emit_event(Evt::Close);
82            }
83            _ => {}
84        }
85        Ok(())
86    }
87}
88
89#[derive(Debug, Clone)]
90pub(crate) enum Evt {
91    FrameSent(Buffer),
92    FrameReceived(Buffer),
93    Error(Value),
94    Close
95}
96
97#[derive(Debug, Clone)]
98pub enum Buffer {
99    Bytes(Vec<u8>),
100    String(String)
101}
102
103impl EventEmitter for WebSocket {
104    type Event = Evt;
105
106    fn tx(&self) -> Option<broadcast::Sender<Self::Event>> { self.tx.lock().unwrap().clone() }
107
108    fn set_tx(&self, tx: broadcast::Sender<Self::Event>) { *self.tx.lock().unwrap() = Some(tx); }
109}
110
111#[derive(Debug, Clone, Copy, PartialEq)]
112pub enum EventType {
113    FrameSent,
114    FrameReceived,
115    Error,
116    Close
117}
118
119impl IsEvent for Evt {
120    type EventType = EventType;
121
122    fn event_type(&self) -> Self::EventType {
123        match self {
124            Evt::FrameSent(_) => EventType::FrameSent,
125            Evt::FrameReceived(_) => EventType::FrameReceived,
126            Evt::Error(_) => EventType::Error,
127            Evt::Close => EventType::Close
128        }
129    }
130}
131
132#[derive(Debug, Deserialize)]
133#[serde(rename_all = "camelCase")]
134struct Initializer {
135    url: String
136}