From 95745f23438f8c902a61a9ff94fc874d68cea815 Mon Sep 17 00:00:00 2001
From: hyh <hongyuanhua@changeself.com>
Date: Thu, 24 Jul 2025 19:08:11 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E6=8E=A5=E6=94=B6=E6=95=B0?=
 =?UTF-8?q?=E6=8D=AE=E8=B6=85=E6=97=B6=E9=97=AE=E9=A2=98=20=E6=B3=A8?=
 =?UTF-8?q?=E5=86=8C=E6=B7=BB=E5=8A=A0=E8=B7=AF=E7=94=B1=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../Extensions/ServiceCollectionExtensions.cs             | 8 ++++++--
 src/X1.WebSocket/Middleware/WebSocketMiddleware.cs        | 4 ++--
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/X1.WebSocket/Extensions/ServiceCollectionExtensions.cs b/src/X1.WebSocket/Extensions/ServiceCollectionExtensions.cs
index 8f13340..09efb5d 100644
--- a/src/X1.WebSocket/Extensions/ServiceCollectionExtensions.cs
+++ b/src/X1.WebSocket/Extensions/ServiceCollectionExtensions.cs
@@ -37,8 +37,12 @@ public static class ServiceCollectionExtensions
         return services;
     }
 
-    public static IApplicationBuilder UseWebSocketMiddleware(this IApplicationBuilder app)
+    public static IApplicationBuilder UseWebSocketMiddleware(this IApplicationBuilder app, string path = "/ws")
     {
-        return app.UseMiddleware<WebSocketMiddleware>();
+        // 路由匹配中间件，仅当路径匹配时才使用 WebSocketMiddleware
+        return app.Map(path, appBuilder =>
+        {
+            appBuilder.UseMiddleware<WebSocketMiddleware>();
+        });
     }
 }
\ No newline at end of file
diff --git a/src/X1.WebSocket/Middleware/WebSocketMiddleware.cs b/src/X1.WebSocket/Middleware/WebSocketMiddleware.cs
index 4190368..3efeb36 100644
--- a/src/X1.WebSocket/Middleware/WebSocketMiddleware.cs
+++ b/src/X1.WebSocket/Middleware/WebSocketMiddleware.cs
@@ -168,7 +168,6 @@ public class WebSocketMiddleware
         var buffer = ArrayPool<byte>.Shared.Rent(1024 * 4);
         // 从对象池获取消息通道
         var messageChannel = _channelPool.Get();
-
         try
         {
             // 创建取消令牌源，设置连接超时
@@ -334,6 +333,7 @@ public class WebSocketMiddleware
                     connectionId, receiveResult.MessageType, receiveResult.Count);
                 await ProcessMessage(webSocket, connectionId, buffer, receiveResult, 
                     messageChannel, messageStartTime, cancellationToken);
+                messageStartTime = DateTime.UtcNow;
             }
             else
             {
@@ -405,7 +405,7 @@ public class WebSocketMiddleware
                     (long)processingTime);
 
                 _messageBuffer.Reset();
-                messageStartTime = DateTime.UtcNow;
+                //messageStartTime = DateTime.UtcNow;
             }
         });