1use std::collections::{HashMap, VecDeque};
14use std::io::ErrorKind;
15use std::sync::atomic::{AtomicU32, Ordering};
16use std::sync::mpsc::{self, Receiver, SyncSender, TryRecvError, TrySendError};
17use std::sync::Arc;
18use std::time::Duration;
19
20use tungstenite::stream::MaybeTlsStream;
21use tungstenite::{Message, WebSocket};
22
23pub const CONNECTING: u32 = 0;
24pub const OPEN: u32 = 1;
25pub const CLOSED: u32 = 2;
26pub const ERROR: u32 = 3;
27
28pub const MAX_CONNECTIONS: usize = 16;
29const QUEUE: usize = 4096;
31const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
32
33struct Conn {
34 state: Arc<AtomicU32>,
35 outgoing: SyncSender<Vec<u8>>,
36 thread: Option<std::thread::JoinHandle<()>>,
37 incoming: Receiver<Vec<u8>>,
38 queue: VecDeque<Vec<u8>>,
39}
40
41#[derive(Clone, Debug, Default)]
44pub struct NetPolicy {
45 pub allowed: bool,
46 pub hosts: Vec<String>,
48}
49
50impl NetPolicy {
51 pub fn new(allowed: bool, hosts: Vec<String>) -> NetPolicy {
52 NetPolicy { allowed, hosts: hosts.into_iter().map(|h| h.to_ascii_lowercase()).collect() }
53 }
54
55 pub fn permits(&self, host: &str) -> bool {
56 if !self.allowed {
57 return false;
58 }
59 let host = host.trim_end_matches('.').to_ascii_lowercase();
60 self.hosts.is_empty()
61 || self.hosts.iter().any(|p| match p.strip_prefix("*.") {
62 Some(domain) => host.len() > domain.len() + 1 && host.ends_with(domain) && host.as_bytes()[host.len() - domain.len() - 1] == b'.',
63 None => host == *p,
64 })
65 }
66
67 pub fn refusal(&self, url: &str) -> Option<String> {
69 if !self.allowed {
70 return Some("run with --allow-net".into());
71 }
72 let host = url_host(url).unwrap_or_default();
73 (!self.permits(&host)).then(|| format!("{host} is not in --allow-net={}", self.hosts.join(",")))
74 }
75}
76
77pub fn url_host(url: &str) -> Option<String> {
79 let rest = url.split_once("://")?.1;
80 let authority = rest.split(['/', '?', '#']).next()?;
81 let authority = authority.rsplit_once('@').map_or(authority, |(_, a)| a);
82 let host = if authority.starts_with('[') {
83 &authority[..=authority.find(']')?]
84 } else {
85 authority.split(':').next()?
86 };
87 (!host.is_empty()).then(|| host.to_ascii_lowercase())
88}
89
90pub struct Net {
91 policy: NetPolicy,
92 conns: HashMap<i32, Conn>,
93 next: i32,
94}
95
96impl Net {
97 pub fn new(allowed: bool) -> Net {
98 Net::with_policy(NetPolicy::new(allowed, Vec::new()))
99 }
100
101 pub fn with_policy(policy: NetPolicy) -> Net {
102 Net { policy, conns: HashMap::new(), next: 1 }
103 }
104
105 pub fn open(&mut self, url: &str) -> i32 {
106 if let Some(why) = self.policy.refusal(url) {
107 eprintln!("[gasm] net: denied connection to {url} ({why})");
108 return -1;
109 }
110 if !url.starts_with("ws://") && !url.starts_with("wss://") {
111 eprintln!("[gasm] net: only ws:// and wss:// URLs are supported: {url}");
112 return -1;
113 }
114 if self.conns.len() >= MAX_CONNECTIONS {
115 eprintln!("[gasm] net: too many connections (max {MAX_CONNECTIONS}): {url}");
116 return -1;
117 }
118 install_crypto_provider();
119 let state = Arc::new(AtomicU32::new(CONNECTING));
120 let (out_tx, out_rx) = mpsc::sync_channel::<Vec<u8>>(QUEUE);
121 let (in_tx, in_rx) = mpsc::sync_channel::<Vec<u8>>(QUEUE);
122 let (st, url_owned) = (state.clone(), url.to_owned());
123 let thread = std::thread::spawn(move || run(&url_owned, &st, out_rx, in_tx));
124 let h = self.next;
125 self.next += 1;
126 self.conns.insert(h, Conn { state, outgoing: out_tx, incoming: in_rx, queue: VecDeque::new(), thread: Some(thread) });
127 h
128 }
129
130 fn check(&self, h: i32) -> Result<bool, String> {
133 if h <= 0 || h >= self.next {
134 return Err(format!("gasm:net: invalid connection handle {h}"));
135 }
136 Ok(self.conns.contains_key(&h))
137 }
138
139 pub fn state(&self, h: i32) -> Result<u32, String> {
140 Ok(match self.check(h)? {
141 true => self.conns[&h].state.load(Ordering::Acquire),
142 false => CLOSED,
143 })
144 }
145
146 pub fn send(&mut self, h: i32, msg: Vec<u8>) -> Result<i32, String> {
148 if !self.check(h)? {
149 return Ok(-1);
150 }
151 let c = &self.conns[&h];
152 if c.state.load(Ordering::Acquire) != OPEN {
153 return Ok(-1);
154 }
155 Ok(match c.outgoing.try_send(msg) {
156 Ok(()) => 0,
157 Err(TrySendError::Full(_)) => {
158 eprintln!("[gasm] net: send queue full ({QUEUE} messages)");
159 -1
160 }
161 Err(TrySendError::Disconnected(_)) => -1,
162 })
163 }
164
165 pub fn peek(&mut self, h: i32) -> Result<Result<&[u8], i32>, String> {
167 if !self.check(h)? {
168 return Ok(Err(-1));
169 }
170 Ok(self.peek_open(h))
171 }
172
173 fn peek_open(&mut self, h: i32) -> Result<&[u8], i32> {
174 let Some(c) = self.conns.get_mut(&h) else { return Err(-1) };
175 while let Ok(m) = c.incoming.try_recv() {
176 c.queue.push_back(m);
177 }
178 if c.queue.is_empty() {
179 return Err(if c.state.load(Ordering::Acquire) >= CLOSED { -1 } else { 0 });
180 }
181 Ok(c.queue.front().unwrap())
182 }
183
184 pub fn pop(&mut self, h: i32) {
185 if let Some(c) = self.conns.get_mut(&h) {
186 c.queue.pop_front();
187 }
188 }
189
190 pub fn close(&mut self, h: i32) -> Result<(), String> {
192 self.check(h)?;
193 self.conns.remove(&h); Ok(())
195 }
196}
197
198impl Drop for Net {
199 fn drop(&mut self) {
202 let threads: Vec<_> = self.conns.drain().filter_map(|(_, mut c)| c.thread.take()).collect();
203 let deadline = std::time::Instant::now() + Duration::from_secs(2);
204 for t in threads {
205 while !t.is_finished() && std::time::Instant::now() < deadline {
206 std::thread::sleep(Duration::from_millis(5));
207 }
208 }
209 }
210}
211
212fn run(url: &str, state: &AtomicU32, outgoing: Receiver<Vec<u8>>, incoming: SyncSender<Vec<u8>>) {
215 let mut ws = match connect(url) {
216 Ok(ws) => ws,
217 Err(e) => {
218 eprintln!("[gasm] net: {url}: {e}");
219 state.store(ERROR, Ordering::Release);
220 return;
221 }
222 };
223 let tcp = match ws.get_mut() {
227 MaybeTlsStream::Plain(s) => Some(&*s),
228 MaybeTlsStream::Rustls(s) => Some(&s.sock),
229 _ => None,
230 };
231 if let Some(tcp) = tcp {
232 let _ = tcp.set_nodelay(true);
233 if let Err(e) = tcp.set_nonblocking(true) {
234 eprintln!("[gasm] net: {url}: {e}");
235 state.store(ERROR, Ordering::Release);
236 return;
237 }
238 }
239 state.store(OPEN, Ordering::Release);
240 let end = pump(&mut ws, &outgoing, &incoming);
241 if let Err(e) = &end {
242 eprintln!("[gasm] net: {url}: {e}");
243 }
244 state.store(if end.is_ok() { CLOSED } else { ERROR }, Ordering::Release);
245 close_gracefully(&mut ws);
246}
247
248fn connect(url: &str) -> Result<WebSocket<MaybeTlsStream<std::net::TcpStream>>, String> {
250 use std::net::ToSocketAddrs;
251 let uri: tungstenite::http::Uri = url.parse().map_err(|e| format!("bad URL: {e}"))?;
252 let host = uri.host().ok_or("URL has no host")?.trim_start_matches('[').trim_end_matches(']');
253 let port = uri.port_u16().unwrap_or(if uri.scheme_str() == Some("wss") { 443 } else { 80 });
254 let mut last = String::from("no address");
255 for addr in (host, port).to_socket_addrs().map_err(|e| e.to_string())? {
256 match std::net::TcpStream::connect_timeout(&addr, CONNECT_TIMEOUT) {
257 Ok(tcp) => {
258 tcp.set_read_timeout(Some(CONNECT_TIMEOUT)).map_err(|e| e.to_string())?;
259 tcp.set_write_timeout(Some(CONNECT_TIMEOUT)).map_err(|e| e.to_string())?;
260 let (ws, _) = tungstenite::client_tls(url, tcp).map_err(|e| e.to_string())?;
261 let tcp = match ws.get_ref() {
262 MaybeTlsStream::Plain(s) => s,
263 MaybeTlsStream::Rustls(s) => &s.sock,
264 _ => return Ok(ws),
265 };
266 tcp.set_read_timeout(None).map_err(|e| e.to_string())?;
267 tcp.set_write_timeout(None).map_err(|e| e.to_string())?;
268 return Ok(ws);
269 }
270 Err(e) => last = format!("{addr}: {e}"),
271 }
272 }
273 Err(last)
274}
275
276fn close_gracefully(ws: &mut WebSocket<MaybeTlsStream<std::net::TcpStream>>) {
280 if ws.close(None).is_err() {
281 return;
282 }
283 let deadline = std::time::Instant::now() + Duration::from_secs(1);
284 while std::time::Instant::now() < deadline {
285 match ws.read() {
286 Ok(_) => {}
287 Err(e) if would_block(&e) => std::thread::sleep(Duration::from_millis(1)),
288 Err(_) => break, }
290 }
291}
292
293pub fn install_crypto_provider() {
295 static ONCE: std::sync::Once = std::sync::Once::new();
296 ONCE.call_once(|| {
297 let _ = rustls::crypto::ring::default_provider().install_default();
298 });
299}
300
301fn would_block(e: &tungstenite::Error) -> bool {
302 matches!(e, tungstenite::Error::Io(e) if matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut | ErrorKind::Interrupted))
303}
304
305pub(crate) fn pump(
309 ws: &mut WebSocket<MaybeTlsStream<std::net::TcpStream>>,
310 outgoing: &Receiver<Vec<u8>>,
311 incoming: &SyncSender<Vec<u8>>,
312) -> Result<(), String> {
313 let mut pending: Option<Vec<u8>> = None;
315 loop {
316 let mut busy = false;
317 loop {
318 match outgoing.try_recv() {
319 Ok(m) => {
320 busy = true;
321 match ws.send(Message::Binary(m.into())) {
322 Ok(()) => {}
323 Err(e) if would_block(&e) => {}
324 Err(e) => return Err(e.to_string()),
325 }
326 }
327 Err(TryRecvError::Empty) => break,
328 Err(TryRecvError::Disconnected) => {
329 flush_all(ws);
330 return Ok(()); }
332 }
333 }
334 match ws.flush() {
335 Ok(()) => {}
336 Err(e) if would_block(&e) => {}
337 Err(e) => return Err(e.to_string()),
338 }
339 if let Some(m) = pending.take() {
340 match incoming.try_send(m) {
341 Ok(()) => busy = true,
342 Err(TrySendError::Full(m)) => pending = Some(m),
343 Err(TrySendError::Disconnected(_)) => return Ok(()),
344 }
345 }
346 let read = if pending.is_some() { Err(tungstenite::Error::Io(ErrorKind::WouldBlock.into())) } else { ws.read() };
347 match read {
348 Ok(Message::Binary(b)) => {
349 busy = true;
350 pending = Some(b.to_vec());
351 }
352 Ok(Message::Text(t)) => {
353 busy = true;
354 pending = Some(t.as_bytes().to_vec());
355 }
356 Ok(Message::Close(_)) => return Ok(()),
357 Ok(_) => busy = true,
358 Err(e) if would_block(&e) => {}
359 Err(tungstenite::Error::ConnectionClosed | tungstenite::Error::AlreadyClosed) => return Ok(()),
360 Err(e) => return Err(e.to_string()),
361 }
362 if !busy {
363 std::thread::sleep(Duration::from_millis(1));
364 }
365 }
366}
367
368fn flush_all(ws: &mut WebSocket<MaybeTlsStream<std::net::TcpStream>>) {
370 let deadline = std::time::Instant::now() + Duration::from_secs(1);
371 while std::time::Instant::now() < deadline {
372 match ws.flush() {
373 Err(e) if would_block(&e) => std::thread::sleep(Duration::from_millis(1)),
374 _ => return,
375 }
376 }
377}
378
379#[cfg(test)]
380mod policy_tests {
381 use super::*;
382
383 #[test]
384 fn hosts() {
385 assert_eq!(url_host("https://User@API.met.no:443/x?y").as_deref(), Some("api.met.no"));
386 assert_eq!(url_host("http://[::1]:8080/").as_deref(), Some("[::1]"));
387 assert_eq!(url_host("ws://localhost").as_deref(), Some("localhost"));
388 let p = NetPolicy::new(true, vec!["api.met.no".into(), "*.example.org".into()]);
389 assert!(p.permits("api.met.no") && p.permits("API.MET.NO.") && p.permits("a.b.example.org"));
390 assert!(!p.permits("met.no") && !p.permits("example.org") && !p.permits("badexample.org"));
391 assert!(NetPolicy::new(true, vec![]).permits("anything"));
392 assert!(!NetPolicy::new(false, vec![]).permits("anything"));
393 }
394}
395
396#[cfg(test)]
397mod tests {
398 use super::*;
399
400 fn settle(net: &Net, h: i32) -> u32 {
402 let t0 = std::time::Instant::now();
403 while net.state(h).unwrap() == CONNECTING && t0.elapsed() < Duration::from_secs(10) {
404 std::thread::sleep(Duration::from_millis(20));
405 }
406 net.state(h).unwrap()
407 }
408
409 #[test]
411 #[ignore = "needs internet"]
412 fn wss_public_echo() {
413 let mut net = Net::new(true);
414 let h = net.open("wss://echo.websocket.org");
415 assert!(h > 0);
416 assert_eq!(settle(&net, h), OPEN, "TLS connection did not open");
417 assert_eq!(net.send(h, b"gasm-tls-check".to_vec()), Ok(0));
418 let t0 = std::time::Instant::now();
419 let mut echoed = false;
420 while t0.elapsed() < Duration::from_secs(10) && !echoed {
421 match net.peek(h).unwrap() {
422 Ok(m) => {
423 echoed = m == b"gasm-tls-check";
424 net.pop(h);
425 }
426 Err(0) => std::thread::sleep(Duration::from_millis(20)),
427 Err(_) => break,
428 }
429 }
430 assert!(echoed, "no echo over wss://");
431 }
432
433 #[test]
434 fn denied_without_permission() {
435 assert_eq!(Net::new(false).open("wss://example.com"), -1);
436 assert_eq!(Net::new(true).open("http://example.com"), -1);
437 }
438
439 #[test]
440 fn unknown_handles_are_errors_closed_ones_are_not() {
441 let mut net = Net::new(true);
442 assert!(net.state(1).is_err() && net.state(0).is_err() && net.close(7).is_err());
443 let h = net.open("ws://127.0.0.1:1"); assert!(h > 0);
445 net.close(h).unwrap();
446 assert_eq!(net.state(h), Ok(CLOSED));
447 assert_eq!(net.send(h, vec![1]), Ok(-1));
448 assert_eq!(net.peek(h).unwrap().err(), Some(-1));
449 net.close(h).unwrap(); }
451}