1use std::net::UdpSocket;
19use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
20use std::sync::{Arc, Mutex};
21
22use thread_relay::{Body, ConnId, Hub, Inbound, Outbound};
23
24const PORT_RANGE: std::ops::Range<u16> = 4200..4210;
28
29pub struct HomeHost {
31 rt: Option<tokio::runtime::Runtime>,
32 port: u16,
33 occupants: Arc<AtomicUsize>,
34}
35
36impl HomeHost {
37 pub fn start(manifest_for: impl FnOnce(u16) -> String) -> Option<HomeHost> {
42 let rt = tokio::runtime::Builder::new_multi_thread()
43 .worker_threads(1)
44 .enable_all()
45 .build()
46 .ok()?;
47 let listener = rt.block_on(async {
48 for port in PORT_RANGE {
49 if let Ok(l) = tokio::net::TcpListener::bind(("0.0.0.0", port)).await {
50 return Some(l);
51 }
52 }
53 None
54 })?;
55 let port = listener.local_addr().ok()?.port();
56 let manifest_text = manifest_for(port);
57 let occupants = Arc::new(AtomicUsize::new(0));
58 rt.spawn(accept_loop(listener, manifest_text, occupants.clone()));
59 tracing::info!("host: serving this place on port {port}");
60 Some(HomeHost {
61 rt: Some(rt),
62 port,
63 occupants,
64 })
65 }
66
67 pub fn port(&self) -> u16 {
68 self.port
69 }
70
71 pub fn occupant_count(&self) -> usize {
74 self.occupants.load(Ordering::Relaxed)
75 }
76
77 pub fn invite_locator(&self) -> String {
79 format!(
80 "thread://{}:{}",
81 lan_ip().unwrap_or_else(|| "127.0.0.1".into()),
82 self.port
83 )
84 }
85
86 pub fn local_room_url(&self) -> String {
88 format!("ws://127.0.0.1:{}/thread/home", self.port)
89 }
90
91 pub fn stop(&mut self) {
93 if let Some(rt) = self.rt.take() {
94 rt.shutdown_background();
95 tracing::info!("host: closed (port {})", self.port);
96 }
97 }
98}
99
100impl Drop for HomeHost {
101 fn drop(&mut self) {
102 self.stop();
103 }
104}
105
106pub fn lan_ip() -> Option<String> {
108 let s = UdpSocket::bind("0.0.0.0:0").ok()?;
109 s.connect("192.0.2.1:9").ok()?; Some(s.local_addr().ok()?.ip().to_string())
111}
112
113async fn accept_loop(
114 listener: tokio::net::TcpListener,
115 manifest_text: String,
116 occupants: Arc<AtomicUsize>,
117) {
118 type Senders =
119 Arc<Mutex<std::collections::HashMap<ConnId, tokio::sync::mpsc::UnboundedSender<Body>>>>;
120 let hub = Arc::new(Mutex::new(Hub::new(64.0, 15)));
121 let senders: Senders = Arc::new(Mutex::new(Default::default()));
122 let next_conn = Arc::new(AtomicU64::new(1));
123 let manifest = Arc::new(manifest_text);
124
125 fn dispatch(hub: &Mutex<Hub>, senders: &Senders, ev: Inbound) {
126 let outs: Vec<Outbound> = hub.lock().unwrap().handle(ev);
127 let s = senders.lock().unwrap();
128 for o in outs {
129 if let Some(tx) = s.get(&o.conn) {
130 let _ = tx.send(o.body);
131 }
132 }
133 }
134
135 loop {
136 let Ok((stream, _)) = listener.accept().await else {
137 break;
138 };
139 let hub = hub.clone();
140 let senders = senders.clone();
141 let manifest = manifest.clone();
142 let conn = next_conn.fetch_add(1, Ordering::Relaxed);
143 let occupants = occupants.clone();
144 tokio::spawn(async move {
145 let mut probe = [0u8; 1024];
147 let Ok(n) = stream.peek(&mut probe).await else {
148 return;
149 };
150 let head = String::from_utf8_lossy(&probe[..n]).to_ascii_lowercase();
151 if !head.contains("upgrade: websocket") {
152 serve_manifest(stream, &manifest).await;
153 return;
154 }
155
156 use futures_util::{SinkExt, StreamExt};
157 use tokio_tungstenite::tungstenite::Message;
158 let Ok(ws) = tokio_tungstenite::accept_async(stream).await else {
159 return;
160 };
161 let (mut write, mut read) = ws.split();
162 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<Body>();
163 senders.lock().unwrap().insert(conn, tx);
164 dispatch(&hub, &senders, Inbound::Connect(conn, "home".to_string()));
166 occupants.fetch_add(1, Ordering::Relaxed);
167
168 loop {
169 tokio::select! {
170 out = rx.recv() => match out {
171 Some(Body::Text(t)) => {
172 if write.send(Message::Text(t.into())).await.is_err() { break; }
173 }
174 Some(Body::Binary(b)) => {
175 if write.send(Message::Binary(b.into())).await.is_err() { break; }
176 }
177 None => break,
178 },
179 msg = read.next() => match msg {
180 Some(Ok(Message::Text(t))) => {
181 dispatch(&hub, &senders, Inbound::Text(conn, t.to_string()))
182 }
183 Some(Ok(Message::Binary(b))) => {
184 dispatch(&hub, &senders, Inbound::Binary(conn, b.to_vec()))
185 }
186 Some(Ok(Message::Close(_))) | None | Some(Err(_)) => break,
187 Some(Ok(_)) => {} },
189 }
190 }
191 occupants.fetch_sub(1, Ordering::Relaxed);
192 senders.lock().unwrap().remove(&conn);
193 dispatch(&hub, &senders, Inbound::Disconnect(conn));
194 });
195 }
196}
197
198async fn serve_manifest(mut stream: tokio::net::TcpStream, manifest: &str) {
202 use tokio::io::{AsyncReadExt, AsyncWriteExt};
203 let mut buf = [0u8; 2048];
205 let _ = stream.read(&mut buf).await;
206 let body = manifest.as_bytes();
207 let head = format!(
208 "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\naccess-control-allow-origin: *\r\ncontent-length: {}\r\nconnection: close\r\n\r\n",
209 body.len()
210 );
211 let _ = stream.write_all(head.as_bytes()).await;
212 let _ = stream.write_all(body).await;
213 let _ = stream.shutdown().await;
214}
215
216#[cfg(test)]
217mod tests {
218 use super::*;
219
220 #[test]
224 fn serves_manifest_and_room_then_kicks_on_stop() {
225 let manifest = r#"{ "thread": "thread/0.1", "world": { "id": "home", "title": "T" } }"#;
226 let mut host = HomeHost::start(|_port| manifest.to_string()).expect("a free invite port");
227 let port = host.port();
228
229 let rt = tokio::runtime::Builder::new_current_thread()
230 .enable_all()
231 .build()
232 .unwrap();
233 rt.block_on(async {
234 use tokio::io::{AsyncReadExt, AsyncWriteExt};
236 let mut s = tokio::net::TcpStream::connect(("127.0.0.1", port))
237 .await
238 .unwrap();
239 s.write_all(b"GET /.well-known/thread/world.json HTTP/1.1\r\nhost: x\r\n\r\n")
240 .await
241 .unwrap();
242 let mut resp = String::new();
243 s.read_to_string(&mut resp).await.unwrap();
244 assert!(resp.starts_with("HTTP/1.1 200"), "manifest served: {resp}");
245 assert!(resp.contains(r#""id": "home""#));
246
247 use futures_util::{SinkExt, StreamExt};
249 use tokio_tungstenite::tungstenite::Message;
250 let url = format!("ws://127.0.0.1:{port}/thread/home");
251 let (mut a, _) = tokio_tungstenite::connect_async(&url).await.unwrap();
252 a.send(Message::Text(r#"{"t":"join"}"#.into()))
253 .await
254 .unwrap();
255 let welcome = a.next().await.unwrap().unwrap();
256 assert!(welcome.to_text().unwrap().contains("welcome"));
257
258 let (mut b, _) = tokio_tungstenite::connect_async(&url).await.unwrap();
259 b.send(Message::Text(r#"{"t":"join"}"#.into()))
260 .await
261 .unwrap();
262 let _welcome_b = b.next().await.unwrap().unwrap();
263 let join_seen_by_a = a.next().await.unwrap().unwrap();
264 assert!(
265 join_seen_by_a.to_text().unwrap().contains("join"),
266 "a sees b arrive"
267 );
268
269 host.stop();
271 let end = tokio::time::timeout(std::time::Duration::from_secs(3), async {
272 loop {
273 match a.next().await {
274 Some(Ok(m)) if m.is_close() => break,
275 None | Some(Err(_)) => break,
276 _ => {}
277 }
278 }
279 })
280 .await;
281 assert!(end.is_ok(), "guest socket closed when the host stopped");
282 });
283 }
284}