1use crate::lua::builtins::json::lua_value_to_json;
6use http_body_util::Full;
7use hyper::body::{Bytes, Frame, Incoming};
8use hyper::server::conn::http1;
9use hyper::service::service_fn;
10use hyper::{Request, Response, StatusCode};
11use mlua::{Lua, Table, Value};
12use std::cell::RefCell;
13use std::collections::HashMap;
14use std::pin::Pin;
15use std::rc::Rc;
16use std::task::{Context, Poll};
17use tokio::net::TcpListener;
18use tracing::error;
19
20#[derive(Clone)]
32pub struct LuaAxumRouter(pub axum::Router);
33
34impl mlua::UserData for LuaAxumRouter {}
35
36pub(super) fn register_serve(lua: &Lua, http_table: &Table) -> mlua::Result<()> {
38 let serve_fn = lua.create_async_function(|lua, args: mlua::MultiValue| async move {
39 let mut args_iter = args.into_iter();
40
41 let port: u16 = match args_iter.next() {
42 Some(Value::Integer(n)) => n as u16,
43 _ => {
44 return Err::<(), _>(mlua::Error::runtime(
45 "http.serve: first argument must be a port number",
46 ));
47 }
48 };
49
50 let routes_table = match args_iter.next() {
51 Some(Value::Table(t)) => t,
52 _ => {
53 return Err::<(), _>(mlua::Error::runtime(
54 "http.serve: second argument must be a routes table",
55 ));
56 }
57 };
58
59 let routes = Rc::new(parse_routes(&routes_table)?);
60
61 let listener = TcpListener::bind(format!("0.0.0.0:{port}"))
62 .await
63 .map_err(|e| mlua::Error::runtime(format!("http.serve: bind failed: {e}")))?;
64
65 let actual_port = listener
67 .local_addr()
68 .map_err(|e| {
69 mlua::Error::runtime(format!("http.serve: failed to get local addr: {e}"))
70 })?
71 .port();
72 lua.globals().set("_SERVER_PORT", actual_port)?;
73
74 loop {
75 let (stream, addr) = listener
76 .accept()
77 .await
78 .map_err(|e| mlua::Error::runtime(format!("http.serve: accept failed: {e}")))?;
79 let peer_addr = addr.to_string();
80
81 let routes = routes.clone();
82 let lua_clone = lua.clone();
83
84 tokio::task::spawn_local(async move {
85 let io = hyper_util::rt::TokioIo::new(stream);
86 let routes = routes.clone();
87 let lua = lua_clone.clone();
88 let peer_addr = peer_addr.clone();
89
90 let service = service_fn(move |req: Request<Incoming>| {
91 let routes = routes.clone();
92 let lua = lua.clone();
93 let peer_addr = peer_addr.clone();
94 async move { handle_request(&lua, &routes, None, peer_addr, req).await }
95 });
96
97 if let Err(e) = http1::Builder::new()
98 .serve_connection(io, service)
99 .with_upgrades()
100 .await
101 && !e.to_string().contains("connection closed")
102 {
103 error!("http.serve: connection error: {e}");
104 }
105 });
106 }
107 })?;
108 http_table.set("serve", serve_fn)?;
109
110 let serve_with_extra_fn =
127 lua.create_async_function(|lua, args: mlua::MultiValue| async move {
128 let mut args_iter = args.into_iter();
129
130 let port: u16 = match args_iter.next() {
131 Some(Value::Integer(n)) => n as u16,
132 _ => {
133 return Err::<(), _>(mlua::Error::runtime(
134 "http.serve_with_extra: first argument must be a port number",
135 ));
136 }
137 };
138
139 let routes_table = match args_iter.next() {
140 Some(Value::Table(t)) => t,
141 _ => {
142 return Err::<(), _>(mlua::Error::runtime(
143 "http.serve_with_extra: second argument must be a routes table",
144 ));
145 }
146 };
147
148 let extra_router: axum::Router = match args_iter.next() {
149 Some(Value::UserData(ud)) => {
150 let r = ud.borrow::<LuaAxumRouter>().map_err(|_| {
151 mlua::Error::runtime(
152 "http.serve_with_extra: third argument must be a LuaAxumRouter userdata",
153 )
154 })?;
155 r.0.clone()
156 }
157 _ => {
158 return Err::<(), _>(mlua::Error::runtime(
159 "http.serve_with_extra: third argument must be a LuaAxumRouter userdata",
160 ));
161 }
162 };
163
164 let routes = Rc::new(parse_routes(&routes_table)?);
165
166 let listener = TcpListener::bind(format!("0.0.0.0:{port}"))
167 .await
168 .map_err(|e| {
169 mlua::Error::runtime(format!("http.serve_with_extra: bind failed: {e}"))
170 })?;
171
172 let actual_port = listener
173 .local_addr()
174 .map_err(|e| {
175 mlua::Error::runtime(format!(
176 "http.serve_with_extra: failed to get local addr: {e}"
177 ))
178 })?
179 .port();
180 lua.globals().set("_SERVER_PORT", actual_port)?;
181
182 loop {
183 let (stream, addr) = listener.accept().await.map_err(|e| {
184 mlua::Error::runtime(format!("http.serve_with_extra: accept failed: {e}"))
185 })?;
186 let peer_addr = addr.to_string();
187
188 let routes = routes.clone();
189 let lua_clone = lua.clone();
190 let extra_router = extra_router.clone();
191
192 tokio::task::spawn_local(async move {
193 let io = hyper_util::rt::TokioIo::new(stream);
194 let routes = routes.clone();
195 let lua = lua_clone.clone();
196 let peer_addr = peer_addr.clone();
197 let extra_router = extra_router.clone();
198
199 let service = service_fn(move |req: Request<Incoming>| {
200 let routes = routes.clone();
201 let lua = lua.clone();
202 let peer_addr = peer_addr.clone();
203 let extra_router = extra_router.clone();
204 async move {
205 handle_request(&lua, &routes, Some(extra_router), peer_addr, req).await
206 }
207 });
208
209 if let Err(e) = http1::Builder::new()
210 .serve_connection(io, service)
211 .with_upgrades()
212 .await
213 && !e.to_string().contains("connection closed")
214 {
215 error!("http.serve_with_extra: connection error: {e}");
216 }
217 });
218 }
219 })?;
220 http_table.set("serve_with_extra", serve_with_extra_fn)?;
221 Ok(())
222}
223
224struct SseBody {
226 rx: tokio::sync::mpsc::Receiver<Bytes>,
227}
228
229impl hyper::body::Body for SseBody {
230 type Data = Bytes;
231 type Error = std::convert::Infallible;
232
233 fn poll_frame(
234 mut self: Pin<&mut Self>,
235 cx: &mut Context<'_>,
236 ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
237 match self.rx.poll_recv(cx) {
238 Poll::Ready(Some(bytes)) => Poll::Ready(Some(Ok(Frame::data(bytes)))),
239 Poll::Ready(None) => Poll::Ready(None),
240 Poll::Pending => Poll::Pending,
241 }
242 }
243}
244
245fn format_sse_event(event_table: &Table) -> mlua::Result<String> {
247 let mut out = String::new();
248
249 if let Ok(Some(event)) = event_table.get::<Option<String>>("event") {
250 if event.contains('\n') || event.contains('\r') {
251 return Err(mlua::Error::runtime(
252 "SSE event name must not contain newlines",
253 ));
254 }
255 out.push_str("event: ");
256 out.push_str(&event);
257 out.push('\n');
258 }
259 if let Ok(Some(data)) = event_table.get::<Option<String>>("data") {
260 for line in data.split('\n') {
262 out.push_str("data: ");
263 out.push_str(line);
264 out.push('\n');
265 }
266 }
267 if let Ok(Some(id)) = event_table.get::<Option<String>>("id") {
268 if id.contains('\n') || id.contains('\r') {
269 return Err(mlua::Error::runtime("SSE id must not contain newlines"));
270 }
271 out.push_str("id: ");
272 out.push_str(&id);
273 out.push('\n');
274 }
275 if let Ok(Some(retry)) = event_table.get::<Option<i64>>("retry") {
276 out.push_str("retry: ");
277 out.push_str(&retry.to_string());
278 out.push('\n');
279 }
280
281 out.push('\n');
283 Ok(out)
284}
285
286fn parse_routes(routes_table: &Table) -> mlua::Result<HashMap<(String, String), mlua::Function>> {
287 let mut routes = HashMap::new();
288 for method_pair in routes_table.pairs::<String, Table>() {
289 let (method, paths_table) = method_pair?;
290 let method_upper = method.to_uppercase();
291 for path_pair in paths_table.pairs::<String, mlua::Function>() {
292 let (path, func) = path_pair?;
293 routes.insert((method_upper.clone(), path), func);
294 }
295 }
296 Ok(routes)
297}
298
299type ServerBody = axum::body::Body;
309
310fn lookup_route<'a>(
311 routes: &'a HashMap<(String, String), mlua::Function>,
312 method: &str,
313 path: &str,
314) -> Option<&'a mlua::Function> {
315 let key = (method.to_string(), path.to_string());
316 if let Some(f) = routes.get(&key) {
317 return Some(f);
318 }
319 let mut search = path;
320 while let Some(pos) = search.rfind('/') {
321 let prefix = &search[..pos];
322 let wildcard_key = (method.to_string(), format!("{prefix}/*"));
323 if let Some(f) = routes.get(&wildcard_key) {
324 return Some(f);
325 }
326 if pos == 0 {
327 let root_key = (method.to_string(), "/*".to_string());
328 return routes.get(&root_key);
329 }
330 search = prefix;
331 }
332 None
333}
334
335fn is_websocket_upgrade(headers: &[(String, String)]) -> bool {
336 headers.iter().any(|(k, v)| {
337 k.eq_ignore_ascii_case("upgrade") && v.to_ascii_lowercase().contains("websocket")
338 })
339}
340
341fn validate_ws_request(headers: &[(String, String)]) -> Result<String, &'static str> {
342 let mut has_connection_upgrade = false;
343 let mut version_ok = false;
344 let mut key: Option<String> = None;
345 for (k, v) in headers {
346 match k.to_ascii_lowercase().as_str() {
347 "connection" if v.to_ascii_lowercase().contains("upgrade") => {
348 has_connection_upgrade = true;
349 }
350 "sec-websocket-version" if v.trim() == "13" => {
351 version_ok = true;
352 }
353 "sec-websocket-key" => {
354 key = Some(v.clone());
355 }
356 _ => {}
357 }
358 }
359 if !has_connection_upgrade {
360 return Err("missing Connection: Upgrade header");
361 }
362 if !version_ok {
363 return Err("Sec-WebSocket-Version must be 13");
364 }
365 key.ok_or("missing Sec-WebSocket-Key header")
366}
367
368fn compute_ws_accept(key: &str) -> String {
369 use sha1::Digest;
370 const MAGIC: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
371 let mut hasher = sha1::Sha1::new();
372 hasher.update(key.as_bytes());
373 hasher.update(MAGIC);
374 let digest = hasher.finalize();
375 data_encoding::BASE64.encode(&digest)
376}
377
378async fn forward_to_axum_router(
390 mut router: axum::Router,
391 req: Request<Incoming>,
392) -> Result<Response<ServerBody>, hyper::Error> {
393 use tower::Service;
394 match <axum::Router as Service<Request<Incoming>>>::call(&mut router, req).await {
399 Ok(resp) => Ok(resp),
400 Err(_) => unreachable!("axum::Router::call is Infallible"),
401 }
402}
403
404async fn handle_request(
405 lua: &Lua,
406 routes: &HashMap<(String, String), mlua::Function>,
407 extra_router: Option<axum::Router>,
408 peer_addr: String,
409 req: Request<Incoming>,
410) -> Result<Response<ServerBody>, hyper::Error> {
411 let method = req.method().to_string();
412 let path = req.uri().path().to_string();
413 let query = req.uri().query().unwrap_or("").to_string();
414 let headers: Vec<(String, String)> = req
415 .headers()
416 .iter()
417 .filter_map(|(k, v)| v.to_str().ok().map(|v| (k.to_string(), v.to_string())))
418 .collect();
419
420 let is_ws = is_websocket_upgrade(&headers);
421
422 let handler = match lookup_route(routes, &method, &path) {
423 Some(h) => h.clone(),
424 None => {
425 if let Some(router) = extra_router {
430 return forward_to_axum_router(router, req).await;
431 }
432 return Ok(Response::builder()
433 .status(StatusCode::NOT_FOUND)
434 .header("content-type", "text/plain")
435 .body(axum::body::Body::new(Full::new(Bytes::from("not found"))))
436 .unwrap());
437 }
438 };
439
440 if is_ws {
441 let lua_resp =
442 match build_lua_request_and_call(lua, &handler, &method, &path, &query, &headers, "")
443 .await
444 {
445 Ok(t) => t,
446 Err(e) => {
447 return Ok(Response::builder()
448 .status(StatusCode::INTERNAL_SERVER_ERROR)
449 .header("content-type", "text/plain")
450 .body(axum::body::Body::new(Full::new(Bytes::from(format!(
451 "handler error: {e}"
452 )))))
453 .unwrap());
454 }
455 };
456
457 if let Ok(Some(ws_fn)) = lua_resp.get::<Option<mlua::Function>>("ws") {
458 return build_ws_upgrade_response(lua, &headers, lua_resp, ws_fn, peer_addr, req);
459 }
460
461 return lua_response_to_http(lua, &lua_resp);
462 }
463
464 let body_bytes = match http_body_util::BodyExt::collect(req.into_body()).await {
465 Ok(collected) => collected.to_bytes(),
466 Err(_) => Bytes::new(),
467 };
468 let body_str = String::from_utf8_lossy(&body_bytes).to_string();
469
470 match build_lua_request_and_call(lua, &handler, &method, &path, &query, &headers, &body_str)
471 .await
472 {
473 Ok(lua_resp) => lua_response_to_http(lua, &lua_resp),
474 Err(e) => Ok(Response::builder()
475 .status(StatusCode::INTERNAL_SERVER_ERROR)
476 .header("content-type", "text/plain")
477 .body(axum::body::Body::new(Full::new(Bytes::from(format!(
478 "handler error: {e}"
479 )))))
480 .unwrap()),
481 }
482}
483
484fn build_ws_upgrade_response(
485 lua: &Lua,
486 headers: &[(String, String)],
487 resp_table: Table,
488 ws_fn: mlua::Function,
489 peer_addr: String,
490 req: Request<Incoming>,
491) -> Result<Response<ServerBody>, hyper::Error> {
492 let key = match validate_ws_request(headers) {
493 Ok(k) => k,
494 Err(msg) => {
495 return Ok(Response::builder()
496 .status(StatusCode::BAD_REQUEST)
497 .header("content-type", "text/plain")
498 .body(axum::body::Body::new(Full::new(Bytes::from(format!(
499 "websocket upgrade rejected: {msg}"
500 )))))
501 .unwrap());
502 }
503 };
504 let accept = compute_ws_accept(&key);
505
506 let mut builder = Response::builder()
507 .status(StatusCode::SWITCHING_PROTOCOLS)
508 .header(hyper::header::UPGRADE, "websocket")
509 .header(hyper::header::CONNECTION, "Upgrade")
510 .header("sec-websocket-accept", accept);
511
512 if let Ok(Some(headers_table)) = resp_table.get::<Option<Table>>("headers") {
515 for pair in headers_table.pairs::<String, mlua::String>().flatten() {
516 let (k, v) = pair;
517 let kl = k.to_ascii_lowercase();
518 if matches!(
519 kl.as_str(),
520 "upgrade" | "connection" | "sec-websocket-accept"
521 ) {
522 continue;
523 }
524 if let Ok(s) = v.to_str() {
525 builder = builder.header(&k, s.as_ref());
526 }
527 }
528 }
529
530 let response = builder
531 .body(axum::body::Body::new(Full::new(Bytes::new())))
532 .unwrap();
533
534 let lua_clone = lua.clone();
535 tokio::task::spawn_local(async move {
536 let upgraded = match hyper::upgrade::on(req).await {
537 Ok(u) => u,
538 Err(e) => {
539 error!("http.serve: ws upgrade failed: {e}");
540 return;
541 }
542 };
543 let io = hyper_util::rt::TokioIo::new(upgraded);
544 let stream = tokio_tungstenite::WebSocketStream::from_raw_socket(
545 io,
546 tokio_tungstenite::tungstenite::protocol::Role::Server,
547 None,
548 )
549 .await;
550 let conn = crate::lua::builtins::ws::WsServerConn::new(stream, peer_addr);
551 let ud = match lua_clone.create_userdata(conn) {
552 Ok(u) => u,
553 Err(e) => {
554 error!("http.serve: ws userdata creation failed: {e}");
555 return;
556 }
557 };
558 if let Err(e) = ws_fn.call_async::<()>(ud).await
559 && !e.to_string().contains("conn:read: ")
560 {
561 error!("http.serve: ws handler error: {e}");
562 }
563 });
564
565 Ok(response)
566}
567
568async fn build_lua_request_and_call(
569 lua: &Lua,
570 handler: &mlua::Function,
571 method: &str,
572 path: &str,
573 query: &str,
574 headers: &[(String, String)],
575 body: &str,
576) -> mlua::Result<Table> {
577 let req_table = lua.create_table()?;
578 req_table.set("method", method.to_string())?;
579 req_table.set("path", path.to_string())?;
580 req_table.set("query", query.to_string())?;
581 req_table.set("body", body.to_string())?;
582
583 let params_table = lua.create_table()?;
588 if !query.is_empty() {
589 for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
590 params_table.set(key.into_owned(), value.into_owned())?;
591 }
592 }
593 req_table.set("params", params_table)?;
594
595 let headers_table = lua.create_table()?;
596 for (k, v) in headers {
597 headers_table.set(k.as_str(), v.as_str())?;
598 }
599 req_table.set("headers", headers_table)?;
600
601 handler.call_async::<Table>(req_table).await
602}
603
604fn lua_response_to_http(
605 lua: &Lua,
606 resp_table: &Table,
607) -> Result<Response<ServerBody>, hyper::Error> {
608 let status = resp_table
609 .get::<Option<u16>>("status")
610 .unwrap_or(None)
611 .unwrap_or(200);
612
613 if let Ok(Some(sse_fn)) = resp_table.get::<Option<mlua::Function>>("sse") {
615 let (tx, rx) = tokio::sync::mpsc::channel::<Bytes>(32);
616
617 let mut builder =
618 Response::builder().status(StatusCode::from_u16(status).unwrap_or(StatusCode::OK));
619
620 if let Ok(Some(headers_table)) = resp_table.get::<Option<Table>>("headers") {
622 for pair in headers_table.pairs::<String, Value>().flatten() {
623 let (k, v) = pair;
624 match v {
625 Value::String(s) => {
626 if let Ok(s) = s.to_str() {
627 builder = builder.header(&k, s.as_ref());
628 }
629 }
630 Value::Table(t) => {
631 for val in t.sequence_values::<String>().flatten() {
634 builder = builder.header(&k, val);
635 }
636 }
637 _ => {}
638 }
639 }
640 }
641
642 let mut response = builder.body(axum::body::Body::new(SseBody { rx })).unwrap();
643 let response_headers = response.headers_mut();
644 if !response_headers.contains_key(hyper::header::CONTENT_TYPE) {
645 response_headers.insert(
646 hyper::header::CONTENT_TYPE,
647 hyper::header::HeaderValue::from_static("text/event-stream"),
648 );
649 }
650 if !response_headers.contains_key(hyper::header::CACHE_CONTROL) {
651 response_headers.insert(
652 hyper::header::CACHE_CONTROL,
653 hyper::header::HeaderValue::from_static("no-cache"),
654 );
655 }
656 if !response_headers.contains_key(hyper::header::CONNECTION) {
657 response_headers.insert(
658 hyper::header::CONNECTION,
659 hyper::header::HeaderValue::from_static("keep-alive"),
660 );
661 }
662
663 let lua_clone = lua.clone();
669 tokio::task::spawn_local(async move {
670 let tx_holder: Rc<RefCell<Option<tokio::sync::mpsc::Sender<Bytes>>>> =
671 Rc::new(RefCell::new(Some(tx)));
672 let tx_for_fn = tx_holder.clone();
673
674 let send_fn = match lua_clone.create_async_function(move |_lua, event_table: Table| {
675 let tx_ref = tx_for_fn.clone();
676 async move {
677 let formatted = format_sse_event(&event_table)?;
678 let tx = tx_ref
679 .borrow()
680 .clone()
681 .ok_or_else(|| mlua::Error::runtime("SSE stream closed"))?;
682 if tx.send(Bytes::from(formatted)).await.is_err() {
683 return Err(mlua::Error::runtime("SSE stream closed"));
684 }
685 Ok(())
686 }
687 }) {
688 Ok(f) => f,
689 Err(e) => {
690 error!("http.serve SSE: failed to create send callback: {e}");
691 return;
692 }
693 };
694
695 if let Err(e) = sse_fn.call_async::<()>(send_fn).await
696 && !e.to_string().contains("SSE stream closed")
697 {
698 error!("http.serve SSE: handler error: {e}");
699 }
700
701 tx_holder.borrow_mut().take();
703 });
704
705 return Ok(response);
706 }
707
708 let mut builder =
709 Response::builder().status(StatusCode::from_u16(status).unwrap_or(StatusCode::OK));
710
711 let has_content_type =
712 if let Ok(Some(headers_table)) = resp_table.get::<Option<Table>>("headers") {
713 let mut found_ct = false;
714 for pair in headers_table.pairs::<String, Value>().flatten() {
715 let (k, v) = pair;
716 if k.eq_ignore_ascii_case("content-type") {
717 found_ct = true;
718 }
719 match v {
720 Value::String(s) => {
721 if let Ok(s) = s.to_str() {
722 builder = builder.header(&k, s.as_ref());
723 }
724 }
725 Value::Table(t) => {
726 for val in t.sequence_values::<String>().flatten() {
729 builder = builder.header(&k, val);
730 }
731 }
732 _ => {}
733 }
734 }
735 found_ct
736 } else {
737 false
738 };
739
740 let body_bytes = if let Ok(Some(json_table)) = resp_table.get::<Option<Table>>("json") {
741 let json_val =
742 lua_value_to_json(&Value::Table(json_table)).unwrap_or(serde_json::Value::Null);
743 let serialized = serde_json::to_string(&json_val).unwrap_or_else(|_| "null".to_string());
744 if !has_content_type {
745 builder = builder.header("content-type", "application/json");
746 }
747 Bytes::from(serialized)
748 } else if let Ok(Some(body_lua)) = resp_table.get::<Option<mlua::String>>("body") {
749 if !has_content_type {
750 builder = builder.header("content-type", "text/plain");
751 }
752 Bytes::from(body_lua.as_bytes().to_vec())
753 } else {
754 if !has_content_type {
755 builder = builder.header("content-type", "text/plain");
756 }
757 Bytes::new()
758 };
759
760 Ok(builder
761 .body(axum::body::Body::new(Full::new(body_bytes)))
762 .unwrap())
763}
764
765#[cfg(all(test, feature = "server"))]
768mod tests {
769 use super::*;
770 use axum::Router;
771 use axum::routing::get;
772 use mlua::Lua;
773
774 #[test]
780 fn lua_axum_router_round_trips_through_mlua_globals() {
781 let lua = Lua::new();
782 let router = Router::new().route("/ping", get(|| async { "pong" }));
783 let wrapped = LuaAxumRouter(router);
784
785 let ud = lua
786 .create_userdata(wrapped)
787 .expect("create_userdata for LuaAxumRouter");
788 lua.globals()
789 .set("EXTRA_ROUTER", ud)
790 .expect("stash userdata in globals");
791
792 let value: mlua::Value = lua
793 .globals()
794 .get("EXTRA_ROUTER")
795 .expect("read userdata back from globals");
796 let ud = match value {
797 mlua::Value::UserData(u) => u,
798 other => panic!("expected UserData, got {other:?}"),
799 };
800 let _borrowed = ud
801 .borrow::<LuaAxumRouter>()
802 .expect("downcast to LuaAxumRouter");
803 }
804
805 #[test]
810 fn lua_axum_router_is_clone_and_preserves_routes() {
811 let router = Router::<()>::new().route("/health", get(|| async { "ok" }));
812 let wrapped = LuaAxumRouter(router);
813 let _cloned = wrapped.clone();
814 }
816}