From 3059929eb92180526dd094bd714baeb1cf221930 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=90=96=E7=8C=A9?= Date: Thu, 4 Jan 2024 16:04:26 +0800 Subject: [PATCH] fix: should break the loop when connection has closed (#18) * fix: should break the loop * chore: modify unit test --------- Co-authored-by: kinggo --- ws_reverse_proxy.go | 5 ++++- ws_reverse_proxy_test.go | 2 ++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/ws_reverse_proxy.go b/ws_reverse_proxy.go index 90aa925..22faffe 100644 --- 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 index b372ea7..83fd506 100644 --- 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())