next_web_websocket/core/
router.rs

1use std::{net::SocketAddr, sync::Arc};
2
3use axum::{
4    extract::{ConnectInfo, State, WebSocketUpgrade},
5    http::Uri,
6    response::IntoResponse,
7    routing::any,
8    Router,
9};
10use next_web_core::{core::router::ApplyRouter, ApplicationContext};
11use rudi_dev::SingleOwner;
12
13use super::{
14    handle_socket::handle_socket, handler::WebSocketHandler, ws_context::WebSocketContext,
15};
16
17#[SingleOwner(binds = [Self::into_router])]
18pub(crate) struct WSRouter;
19
20impl WSRouter {
21    fn into_router(self) -> Box<dyn ApplyRouter> {
22        Box::new(self)
23    }
24}
25
26impl ApplyRouter for WSRouter {
27    fn router(&self, ctx: &mut ApplicationContext) -> axum::Router {
28        let mut router = Router::new();
29        let mut context = ctx.resolve::<WebSocketContext>();
30        let handlers = ctx.resolve_by_type::<Arc<dyn WebSocketHandler>>();
31
32        for item in handlers.iter() {
33            for path in item.paths() {
34                router = router.route(path, any(ws_handler));
35                context.add_handler(path, item.clone());
36            }
37        }
38        router.with_state(Arc::new(context))
39    }
40}
41
42/// The handler for the HTTP request (this gets called when the HTTP request lands at the start
43/// of websocket negotiation). After this completes, the actual switching from HTTP to
44/// websocket protocol will occur.
45/// This is the last point where we can extract TCP/IP metadata such as IP address of the client
46/// as well as things from HTTP headers such as user-agent of the browser etc.
47async fn ws_handler(
48    ws: WebSocketUpgrade,
49    ConnectInfo(addr): ConnectInfo<SocketAddr>,
50    State(ctx): State<Arc<WebSocketContext>>,
51    uri: Uri,
52) -> impl IntoResponse {
53    // finalize the upgrade process by returning upgrade callback.
54    // we can customize the callback by sending additional info such as address.
55    let properties = ctx.properties();
56    let path = uri.path().to_string();
57    ws.max_message_size(properties.max_msg_size().unwrap_or(64 << 20))
58        .max_write_buffer_size(properties.max_write_buffer_size().unwrap_or(usize::MAX))
59        .on_upgrade(move |socket| handle_socket(socket, ctx, addr, path))
60}
next_web_websocket/core/router.rs

next_web_websocket/core/
router.rs
