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

next_web_websocket/core/
router.rs
