From 3a99b1fb247b0e4fb8fdc0e9681cad8624dafdc4 Mon Sep 17 00:00:00 2001
From: Chrislearn Young <chris@acroidea.com>
Date: Sat, 28 Sep 2024 11:03:09 +0800
Subject: [PATCH] fix: Accepting a Browser Initiated WebTransport Session
 Always Fails

---
 crates/core/src/conn/quinn/builder.rs |  7 ++++---
 crates/core/src/http/request.rs       | 30 +++++++++++++++++++--------
 crates/core/src/service.rs            | 20 +++++++++---------
 3 files changed, 35 insertions(+), 22 deletions(-)

diff --git a/crates/core/src/conn/quinn/builder.rs b/crates/core/src/conn/quinn/builder.rs
index 06a7102bc..61699ffdb 100644
--- a/crates/core/src/conn/quinn/builder.rs
+++ b/crates/core/src/conn/quinn/builder.rs
@@ -149,15 +149,16 @@ async fn process_web_transport(
     } else {
         conn = response
             .extensions_mut()
-            .remove::<Mutex<salvo_http3::server::Connection<salvo_http3::http3_quinn::Connection, Bytes>>>()
+            .remove::<Arc<Mutex<salvo_http3::server::Connection<salvo_http3::http3_quinn::Connection, Bytes>>>>()
             .map(|c| {
-                c.into_inner()
+                Arc::into_inner(c).unwrap().into_inner()
                     .map_err(|e| IoError::new(ErrorKind::Other, format!("failed to get conn : {}", e)))
             })
             .transpose()?;
         stream = response
             .extensions_mut()
-            .remove::<salvo_http3::server::RequestStream<salvo_http3::http3_quinn::BidiStream<Bytes>, Bytes>>();
+            .remove::<Arc<salvo_http3::server::RequestStream<salvo_http3::http3_quinn::BidiStream<Bytes>, Bytes>>>()
+            .and_then(|stream|Arc::into_inner(stream));
     }
 
     let Some(conn) = conn else {
diff --git a/crates/core/src/http/request.rs b/crates/core/src/http/request.rs
index 1887adf7a..f2519803b 100644
--- a/crates/core/src/http/request.rs
+++ b/crates/core/src/http/request.rs
@@ -527,20 +527,32 @@ impl Request {
         pub async fn web_transport_mut(&mut self) -> Result<&mut crate::proto::WebTransportSession<salvo_http3::http3_quinn::Connection, Bytes>, crate::Error> {
             if self.is_wt_connect() {
                 if self.extensions.get::<crate::proto::WebTransportSession<salvo_http3::http3_quinn::Connection, Bytes>>().is_none() {
-                    let conn = self.extensions.remove::<std::sync::Mutex<salvo_http3::server::Connection<salvo_http3::http3_quinn::Connection, Bytes>>>();
-                    let stream = self.extensions.remove::<salvo_http3::server::RequestStream<salvo_http3::http3_quinn::BidiStream<Bytes>, Bytes>>();
+                    let conn = self.extensions.remove::<Arc<std::sync::Mutex<salvo_http3::server::Connection<salvo_http3::http3_quinn::Connection, Bytes>>>>();
+                    let stream = self.extensions.remove::<Arc<salvo_http3::server::RequestStream<salvo_http3::http3_quinn::BidiStream<Bytes>, Bytes>>>();
                     match (conn, stream) {
                         (Some(conn), Some(stream)) => {
-                            if let Ok(conn) = conn.into_inner() {
-                                    let session =  crate::proto::WebTransportSession::accept(stream, conn).await?;
-                                    self.extensions.insert(Arc::new(session));
-                                    if let Some(session) = self.extensions.get_mut::<crate::proto::WebTransportSession<salvo_http3::http3_quinn::Connection, Bytes>>() {
-                                        Ok(session)
+                            if let Some(conn) = Arc::into_inner(conn) {
+                                if let Ok(conn) = conn.into_inner() {
+                                    if let Some(stream) = Arc::into_inner(stream) {
+                                        let session =  crate::proto::WebTransportSession::accept(stream, conn).await?;
+                                        self.extensions.insert(Arc::new(session));
+                                        if let Some(session) = self.extensions.get_mut::<Arc<crate::proto::WebTransportSession<salvo_http3::http3_quinn::Connection, Bytes>>>() {
+                                            if let Some(session) = Arc::get_mut(session) {
+                                                Ok(session)
+                                            } else {
+                                                Err(crate::Error::Other("web transport session should not used twice".into()))
+                                            }
+                                        } else {
+                                            Err(crate::Error::Other("web transport session not found in request extension".into()))
+                                        }
                                     } else {
-                                        Err(crate::Error::Other("invalid web transport".into()))
+                                        Err(crate::Error::Other("web transport stream should not used twice".into()))
                                     }
+                                } else {
+                                    Err(crate::Error::Other("invalid web transport".into()))
+                                }
                             } else {
-                                Err(crate::Error::Other("invalid web transport".into()))
+                                Err(crate::Error::Other("quinn connection should not used twice".into()))
                             }
                         }
                         (Some(conn), None) => {
diff --git a/crates/core/src/service.rs b/crates/core/src/service.rs
index 0ff13394f..d01a4025d 100644
--- a/crates/core/src/service.rs
+++ b/crates/core/src/service.rs
@@ -318,25 +318,25 @@ impl HyperHandler {
             #[cfg(feature = "quinn")]
             {
                 use bytes::Bytes;
-                use parking_lot::Mutex;
+                use std::sync::Mutex;
                 if let Some(session) =
-                    req.extensions.remove::<crate::proto::WebTransportSession<
+                    req.extensions.remove::<Arc<crate::proto::WebTransportSession<
                         salvo_http3::http3_quinn::Connection,
                         Bytes,
-                    >>()
+                    >>>()
                 {
-                    res.extensions.insert(Arc::new(session));
+                    res.extensions.insert(session);
                 }
-                if let Some(conn) = req.extensions.remove::<Mutex<
+                if let Some(conn) = req.extensions.remove::<Arc<Mutex<
                     salvo_http3::server::Connection<salvo_http3::http3_quinn::Connection, Bytes>,
-                >>() {
-                    res.extensions.insert(Arc::new(conn));
+                >>>() {
+                    res.extensions.insert(conn);
                 }
-                if let Some(stream) = req.extensions.remove::<salvo_http3::server::RequestStream<
+                if let Some(stream) = req.extensions.remove::<Arc<salvo_http3::server::RequestStream<
                     salvo_http3::http3_quinn::BidiStream<Bytes>,
                     Bytes,
-                >>() {
-                    res.extensions.insert(Arc::new(stream));
+                >>>() {
+                    res.extensions.insert(stream);
                 }
             }
             res