1use std::future::Future;
4use std::sync::Arc;
5
6use sova_core::{App, Request, Router};
7
8use crate::upgrade::upgrade_ws;
9use crate::WsSession;
10
11pub trait WsRouteExt {
13 fn ws<F, Fut>(&mut self, path: &str, handler: F) -> &mut Self
14 where
15 F: Fn(WsSession) -> Fut + Clone + Send + Sync + 'static,
16 Fut: Future<Output = ()> + Send + 'static;
17}
18
19impl WsRouteExt for Router {
20 fn ws<F, Fut>(&mut self, path: &str, handler: F) -> &mut Self
21 where
22 F: Fn(WsSession) -> Fut + Clone + Send + Sync + 'static,
23 Fut: Future<Output = ()> + Send + 'static,
24 {
25 let handler = Arc::new(handler);
26 self.get(path, move |req: Request| {
27 let handler = Arc::clone(&handler);
28 async move {
29 match upgrade_ws(req, move |session| handler(session)).await {
30 Ok(res) => res,
31 Err(res) => res,
32 }
33 }
34 });
35 self
36 }
37}
38
39impl WsRouteExt for App {
40 fn ws<F, Fut>(&mut self, path: &str, handler: F) -> &mut Self
41 where
42 F: Fn(WsSession) -> Fut + Clone + Send + Sync + 'static,
43 Fut: Future<Output = ()> + Send + 'static,
44 {
45 Router::ws(self, path, handler);
46 self
47 }
48}