From 95d7b101a0130b0183ce8b06b9ec44cd6723e847 Mon Sep 17 00:00:00 2001
From: Toshit Chawda <r58playz@gmail.com>
Date: Fri, 27 Sep 2024 16:11:13 -0700
Subject: [PATCH] fix twisp extension

---
 server/src/handle/wisp/mod.rs   | 2 +-
 server/src/handle/wisp/twisp.rs | 6 +++---
 wisp/src/lib.rs                 | 5 +++++
 3 files changed, 9 insertions(+), 4 deletions(-)

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<dyn ProtocolExtensionBuilder + Send + Sync> {
-	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.