1use std::collections::{HashMap, VecDeque};
9use std::io::ErrorKind;
10use std::sync::atomic::{AtomicU32, Ordering};
11use std::sync::mpsc::{self, Receiver, Sender, TryRecvError};
12use std::sync::Arc;
13use std::time::Duration;
14
15use tungstenite::stream::MaybeTlsStream;
16use tungstenite::{Message, WebSocket};
17
18pub const CONNECTING: u32 = 0;
19pub const OPEN: u32 = 1;
20pub const CLOSED: u32 = 2;
21pub const ERROR: u32 = 3;
22
23struct Conn {
24 state: Arc<AtomicU32>,
25 outgoing: Sender<Vec<u8>>,
26 thread: Option<std::thread::JoinHandle<()>>,
27 incoming: Receiver<Vec<u8>>,
28 queue: VecDeque<Vec<u8>>,
29}
30
31pub struct Net {
32 allowed: bool,
33 conns: HashMap<i32, Conn>,
34 next: i32,
35}
36
37impl Net {
38 pub fn new(allowed: bool) -> Net {
39 Net { allowed, conns: HashMap::new(), next: 1 }
40 }
41
42 pub fn open(&mut self, url: &str) -> i32 {
43 if !self.allowed {
44 eprintln!("[gasm] net: denied connection to {url} (run with --allow-net)");
45 return -1;
46 }
47 if !url.starts_with("ws://") && !url.starts_with("wss://") {
48 eprintln!("[gasm] net: only ws:// and wss:// URLs are supported: {url}");
49 return -1;
50 }
51 install_crypto_provider();
52 let state = Arc::new(AtomicU32::new(CONNECTING));
53 let (out_tx, out_rx) = mpsc::channel::<Vec<u8>>();
54 let (in_tx, in_rx) = mpsc::channel::<Vec<u8>>();
55 let (st, url_owned) = (state.clone(), url.to_owned());
56 let thread = std::thread::spawn(move || run(&url_owned, &st, out_rx, in_tx));
57 let h = self.next;
58 self.next += 1;
59 self.conns.insert(h, Conn { state, outgoing: out_tx, incoming: in_rx, queue: VecDeque::new(), thread: Some(thread) });
60 h
61 }
62
63 pub fn state(&self, h: i32) -> u32 {
64 self.conns.get(&h).map_or(ERROR, |c| c.state.load(Ordering::Acquire))
65 }
66
67 pub fn send(&mut self, h: i32, msg: Vec<u8>) -> i32 {
68 match self.conns.get(&h) {
69 Some(c) if c.state.load(Ordering::Acquire) == OPEN && c.outgoing.send(msg).is_ok() => 0,
70 _ => -1,
71 }
72 }
73
74 pub fn peek(&mut self, h: i32) -> Result<&[u8], i32> {
76 let Some(c) = self.conns.get_mut(&h) else { return Err(-1) };
77 while let Ok(m) = c.incoming.try_recv() {
78 c.queue.push_back(m);
79 }
80 if c.queue.is_empty() {
81 return Err(if c.state.load(Ordering::Acquire) >= CLOSED { -1 } else { 0 });
82 }
83 Ok(c.queue.front().unwrap())
84 }
85
86 pub fn pop(&mut self, h: i32) {
87 if let Some(c) = self.conns.get_mut(&h) {
88 c.queue.pop_front();
89 }
90 }
91
92 pub fn close(&mut self, h: i32) {
93 self.conns.remove(&h); }
95}
96
97impl Drop for Net {
98 fn drop(&mut self) {
101 let threads: Vec<_> = self.conns.drain().filter_map(|(_, mut c)| c.thread.take()).collect();
102 let deadline = std::time::Instant::now() + Duration::from_secs(2);
103 for t in threads {
104 while !t.is_finished() && std::time::Instant::now() < deadline {
105 std::thread::sleep(Duration::from_millis(5));
106 }
107 }
108 }
109}
110
111fn run(url: &str, state: &AtomicU32, outgoing: Receiver<Vec<u8>>, incoming: Sender<Vec<u8>>) {
114 let mut ws = match tungstenite::connect(url) {
115 Ok((ws, _)) => ws,
116 Err(e) => {
117 eprintln!("[gasm] net: {url}: {e}");
118 state.store(ERROR, Ordering::Release);
119 return;
120 }
121 };
122 let tcp = match ws.get_mut() {
126 MaybeTlsStream::Plain(s) => Some(&*s),
127 MaybeTlsStream::Rustls(s) => Some(&s.sock),
128 _ => None,
129 };
130 if let Some(tcp) = tcp {
131 let _ = tcp.set_nodelay(true);
132 if let Err(e) = tcp.set_nonblocking(true) {
133 eprintln!("[gasm] net: {url}: {e}");
134 state.store(ERROR, Ordering::Release);
135 return;
136 }
137 }
138 state.store(OPEN, Ordering::Release);
139 let end = pump(&mut ws, &outgoing, &incoming);
140 if let Err(e) = &end {
141 eprintln!("[gasm] net: {url}: {e}");
142 }
143 state.store(if end.is_ok() { CLOSED } else { ERROR }, Ordering::Release);
144 close_gracefully(&mut ws);
145}
146
147fn close_gracefully(ws: &mut WebSocket<MaybeTlsStream<std::net::TcpStream>>) {
151 if ws.close(None).is_err() {
152 return;
153 }
154 let deadline = std::time::Instant::now() + Duration::from_secs(1);
155 while std::time::Instant::now() < deadline {
156 match ws.read() {
157 Ok(_) => {}
158 Err(e) if would_block(&e) => std::thread::sleep(Duration::from_millis(1)),
159 Err(_) => break, }
161 }
162}
163
164pub fn install_crypto_provider() {
166 static ONCE: std::sync::Once = std::sync::Once::new();
167 ONCE.call_once(|| {
168 let _ = rustls::crypto::ring::default_provider().install_default();
169 });
170}
171
172fn would_block(e: &tungstenite::Error) -> bool {
173 matches!(e, tungstenite::Error::Io(e) if matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut | ErrorKind::Interrupted))
174}
175
176pub(crate) fn pump(
180 ws: &mut WebSocket<MaybeTlsStream<std::net::TcpStream>>,
181 outgoing: &Receiver<Vec<u8>>,
182 incoming: &Sender<Vec<u8>>,
183) -> Result<(), String> {
184 loop {
185 let mut busy = false;
186 loop {
187 match outgoing.try_recv() {
188 Ok(m) => {
189 busy = true;
190 match ws.send(Message::Binary(m.into())) {
191 Ok(()) => {}
192 Err(e) if would_block(&e) => {}
193 Err(e) => return Err(e.to_string()),
194 }
195 }
196 Err(TryRecvError::Empty) => break,
197 Err(TryRecvError::Disconnected) => {
198 flush_all(ws);
199 return Ok(()); }
201 }
202 }
203 match ws.flush() {
204 Ok(()) => {}
205 Err(e) if would_block(&e) => {}
206 Err(e) => return Err(e.to_string()),
207 }
208 match ws.read() {
209 Ok(Message::Binary(b)) => {
210 busy = true;
211 if incoming.send(b.to_vec()).is_err() {
212 return Ok(());
213 }
214 }
215 Ok(Message::Text(t)) => {
216 busy = true;
217 let _ = incoming.send(t.as_bytes().to_vec());
218 }
219 Ok(Message::Close(_)) => return Ok(()),
220 Ok(_) => busy = true,
221 Err(e) if would_block(&e) => {}
222 Err(tungstenite::Error::ConnectionClosed | tungstenite::Error::AlreadyClosed) => return Ok(()),
223 Err(e) => return Err(e.to_string()),
224 }
225 if !busy {
226 std::thread::sleep(Duration::from_millis(1));
227 }
228 }
229}
230
231fn flush_all(ws: &mut WebSocket<MaybeTlsStream<std::net::TcpStream>>) {
233 let deadline = std::time::Instant::now() + Duration::from_secs(1);
234 while std::time::Instant::now() < deadline {
235 match ws.flush() {
236 Err(e) if would_block(&e) => std::thread::sleep(Duration::from_millis(1)),
237 _ => return,
238 }
239 }
240}
241
242#[cfg(test)]
243mod tests {
244 use super::*;
245
246 fn settle(net: &Net, h: i32) -> u32 {
248 let t0 = std::time::Instant::now();
249 while net.state(h) == CONNECTING && t0.elapsed() < Duration::from_secs(10) {
250 std::thread::sleep(Duration::from_millis(20));
251 }
252 net.state(h)
253 }
254
255 #[test]
257 #[ignore = "needs internet"]
258 fn wss_public_echo() {
259 let mut net = Net::new(true);
260 let h = net.open("wss://echo.websocket.org");
261 assert!(h > 0);
262 assert_eq!(settle(&net, h), OPEN, "TLS connection did not open");
263 assert_eq!(net.send(h, b"gasm-tls-check".to_vec()), 0);
264 let t0 = std::time::Instant::now();
265 let mut echoed = false;
266 while t0.elapsed() < Duration::from_secs(10) && !echoed {
267 match net.peek(h) {
268 Ok(m) => {
269 echoed = m == b"gasm-tls-check";
270 net.pop(h);
271 }
272 Err(0) => std::thread::sleep(Duration::from_millis(20)),
273 Err(_) => break,
274 }
275 }
276 assert!(echoed, "no echo over wss://");
277 }
278
279 #[test]
280 fn denied_without_permission() {
281 assert_eq!(Net::new(false).open("wss://example.com"), -1);
282 assert_eq!(Net::new(true).open("http://example.com"), -1);
283 }
284}