redevplugin_worker_sdk/
websocket.rs

1use crate::api;
2use crate::error::{Error, Result};
3use crate::http::Header;
4use crate::resource::{
5    Handle, IO_FLAG_BINARY, IO_FLAG_EOF, IO_FLAG_MESSAGE_END, IO_FLAG_TEXT, MAX_IO_CHUNK_BYTES,
6};
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, Serialize)]
10pub struct WebSocketOpen {
11    /// Connector declared by the current method; no implicit default is used.
12    pub connector_id: String,
13    pub url: String,
14    #[serde(default)]
15    pub headers: Vec<Header>,
16    #[serde(default)]
17    pub subprotocols: Vec<String>,
18    #[serde(default)]
19    #[serde(skip_serializing_if = "Option::is_none")]
20    pub timeout_ms: Option<u32>,
21}
22
23pub type OpenOptions = WebSocketOpen;
24
25#[derive(Deserialize)]
26struct OpenResult {
27    handle: u64,
28    #[serde(default)]
29    protocol: String,
30    #[serde(default)]
31    response_headers: Vec<Header>,
32}
33
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub enum Message {
36    Text(String),
37    Binary(Vec<u8>),
38}
39
40pub struct WebSocket {
41    handle: Handle,
42    pub protocol: String,
43    pub response_headers: Vec<Header>,
44}
45
46impl WebSocket {
47    pub fn open(options: WebSocketOpen) -> Result<Self> {
48        let opened: OpenResult = api::call("net.websocket.open", &options)?;
49        Ok(Self {
50            handle: Handle::new(opened.handle)?,
51            protocol: opened.protocol,
52            response_headers: opened.response_headers,
53        })
54    }
55
56    pub fn send_text(&mut self, text: &str) -> Result<()> {
57        self.send_message(text.as_bytes(), IO_FLAG_TEXT)
58    }
59
60    pub fn send_binary(&mut self, bytes: &[u8]) -> Result<()> {
61        self.send_message(bytes, IO_FLAG_BINARY)
62    }
63
64    fn send_message(&mut self, bytes: &[u8], kind: u32) -> Result<()> {
65        if bytes.is_empty() {
66            return self.handle.write(bytes, kind | IO_FLAG_MESSAGE_END);
67        }
68        let chunk_count = bytes.len().div_ceil(MAX_IO_CHUNK_BYTES);
69        for (index, chunk) in bytes.chunks(MAX_IO_CHUNK_BYTES).enumerate() {
70            let mut flags = if index == 0 { kind } else { 0 };
71            if index + 1 == chunk_count {
72                flags |= IO_FLAG_MESSAGE_END;
73            }
74            self.handle.write(chunk, flags)?;
75        }
76        Ok(())
77    }
78
79    pub fn receive(&mut self) -> Result<Message> {
80        let mut body = Vec::new();
81        let mut kind = 0;
82        loop {
83            let (chunk, flags) = self.handle.read(MAX_IO_CHUNK_BYTES)?;
84            let message_kind = flags & (IO_FLAG_TEXT | IO_FLAG_BINARY);
85            if flags & IO_FLAG_EOF != 0 {
86                return Err(Error::from_abi_status(-5));
87            }
88            if message_kind == IO_FLAG_TEXT | IO_FLAG_BINARY {
89                return Err(Error::internal("WebSocket message type is invalid"));
90            }
91            if message_kind != 0 {
92                if kind != 0 {
93                    return Err(Error::internal(
94                        "WebSocket message type changed mid-message",
95                    ));
96                }
97                kind = message_kind;
98            }
99            body.extend_from_slice(&chunk);
100            if flags & IO_FLAG_MESSAGE_END != 0 {
101                return match kind {
102                    IO_FLAG_TEXT => String::from_utf8(body)
103                        .map(Message::Text)
104                        .map_err(|_| Error::internal("WebSocket text message is not UTF-8")),
105                    IO_FLAG_BINARY => Ok(Message::Binary(body)),
106                    _ => Err(Error::internal("WebSocket message type is missing")),
107                };
108            }
109            if chunk.is_empty() {
110                return Err(Error::internal("WebSocket receive made no progress"));
111            }
112        }
113    }
114
115    pub fn ping(&mut self) -> Result<()> {
116        #[derive(Serialize)]
117        struct Arguments {
118            handle: u64,
119        }
120        let _: serde_json::Value = api::call(
121            "net.websocket.ping",
122            &Arguments {
123                handle: self.handle.id(),
124            },
125        )?;
126        Ok(())
127    }
128
129    pub fn close(mut self, code: u16, reason: &str) -> Result<()> {
130        #[derive(Serialize)]
131        struct Arguments<'a> {
132            handle: u64,
133            code: u16,
134            reason: &'a str,
135        }
136        let _: serde_json::Value = api::call(
137            "net.websocket.close",
138            &Arguments {
139                handle: self.handle.id(),
140                code,
141                reason,
142            },
143        )?;
144        self.handle.disarm();
145        Ok(())
146    }
147}
redevplugin_worker_sdk/websocket.rs

redevplugin_worker_sdk/
websocket.rs
