diff --git a/server/src/handle/wisp/mod.rs b/server/src/handle/wisp/mod.rs index f5d038d..11687ec 100644 --- a/server/src/handle/wisp/mod.rs +++ b/server/src/handle/wisp/mod.rs @@ -221,7 +221,7 @@ pub async fn handle_wisp(stream: WispResult, id: String) -> anyhow::Result<()> { let extensions = match extensions { Some(mut exts) => { - exts.push(twisp::new_ext(twisp_map.clone())); + exts.add_extension(twisp::new_ext(twisp_map.clone())); Some(exts) }, None => { diff --git a/server/src/handle/wisp/twisp.rs b/server/src/handle/wisp/twisp.rs index cbf0071..5a9e3e6 100644 --- a/server/src/handle/wisp/twisp.rs +++ b/server/src/handle/wisp/twisp.rs @@ -10,7 +10,7 @@ use pty_process::{Pty, Size}; use tokio::{io::copy, process::Child, select, sync::Mutex}; use tokio_util::compat::{FuturesAsyncReadCompatExt, FuturesAsyncWriteCompatExt}; use wisp_mux::{ - extensions::{AnyProtocolExtension, ProtocolExtension, ProtocolExtensionBuilder}, + extensions::{AnyProtocolExtension, ProtocolExtension, ProtocolExtensionBuilder, AnyProtocolExtensionBuilder}, ws::{LockedWebSocketWrite, WebSocketRead}, MuxStreamAsyncRead, MuxStreamAsyncWrite, WispError, }; @@ -120,8 +120,8 @@ pub fn new_map() -> TwispMap { Arc::new(Mutex::new(HashMap::new())) } -pub fn new_ext(map: TwispMap) -> Box { - Box::new(TWispServerProtocolExtensionBuilder(map)) +pub fn new_ext(map: TwispMap) -> AnyProtocolExtensionBuilder { + AnyProtocolExtensionBuilder::new(TWispServerProtocolExtensionBuilder(map)) } pub async fn handle_twisp( diff --git a/wisp/src/lib.rs b/wisp/src/lib.rs index 53a527c..f5859f1 100644 --- a/wisp/src/lib.rs +++ b/wisp/src/lib.rs @@ -263,6 +263,11 @@ impl WispV2Extensions { closure: Box::new(closure), } } + + /// Add a Wisp V2 extension builder to the settings struct. + pub fn add_extension(&mut self, extension: AnyProtocolExtensionBuilder) { + self.builders.push(extension); + } } /// Server-side multiplexor.