fix: should break the loop when connection has closed (#18)

* fix: should break the loop * chore: modify unit test --------- Co-authored-by: kinggo <[email protected]>
hertz-contrib · Jan 4, 2024 · 3059929 · 3059929
1 parent f728214
commit 3059929
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/ws_reverse_proxy.go b/ws_reverse_proxy.go
@@ -137,9 +137,12 @@ func (w *WSReverseProxy) ServeHTTP(ctx context.Context, c *app.RequestContext) {
 
 			var ce *websocket.CloseError
 			var hzce *hzws.CloseError
-			if !errors.As(err, &ce) || !errors.As(err, &hzce) {
+			if !errors.As(err, &ce) && !errors.As(err, &hzce) {
 				hlog.CtxErrorf(ctx, errMsg, err)
+				continue
 			}
+
+			break
 		}
 	}); err != nil {
 		hlog.CtxErrorf(ctx, "can not upgrade to websocket: %v", err)

diff --git a/ws_reverse_proxy_test.go b/ws_reverse_proxy_test.go
@@ -55,6 +55,7 @@ func TestProxy(t *testing.T) {
 
 	// proxy server
 	ps := server.Default(server.WithHostPorts(":7777"))
+	ps.NoHijackConnPool = true
 	ps.GET("/proxy", proxy.ServeHTTP)
 	go ps.Spin()
 
@@ -63,6 +64,7 @@ func TestProxy(t *testing.T) {
 	go func() {
 		// backend server
 		bs := server.Default()
+		bs.NoHijackConnPool = true
 		bs.GET("/", func(ctx context.Context, c *app.RequestContext) {
 			// Don't upgrade if original host header isn't preserved
 			host := string(c.Host())