closing handshake for uni_websocket

Author: FZambia

internal/configtypes/types.go

@@ -121,7 +121,15 @@ type UniWebSocket struct {
 	WriteBufferSize    int      `mapstructure:"write_buffer_size" json:"write_buffer_size" envconfig:"write_buffer_size" yaml:"write_buffer_size" toml:"write_buffer_size"`
 	WriteTimeout       Duration `mapstructure:"write_timeout" json:"write_timeout" envconfig:"write_timeout" default:"1000ms" yaml:"write_timeout" toml:"write_timeout"`
 	MessageSizeLimit   int      `mapstructure:"message_size_limit" json:"message_size_limit" envconfig:"message_size_limit" default:"65536" yaml:"message_size_limit" toml:"message_size_limit"`
-
+	// DisableClosingHandshake disables WebSocket closing handshake. This restores the behavior prior to
+	// Centrifugo v6.5.1 where server never sent a close frame on connection close initiated by server.
+	// Normally closing handshake is recommended to be performed according to WebSocket protocol RFC,
+	// so this option is useful only in some specific cases when you need to restore the previous behavior.
+	DisableClosingHandshake bool `mapstructure:"disable_closing_handshake" json:"disable_closing_handshake" envconfig:"disable_closing_handshake" yaml:"disable_closing_handshake" toml:"disable_closing_handshake"`
+	// DisableDisconnectPush disables sending disconnect push messages to clients. It's sent by default to make
+	// unidirectional transports similar, but since uni_websocket transport also sends close frame to the client
+	// with the same code/reason – some users may want to disable disconnect push to avoid ambiguity.
+	DisableDisconnectPush bool `mapstructure:"disable_disconnect_push" json:"disable_disconnect_push" envconfig:"disable_disconnect_push" yaml:"disable_disconnect_push" toml:"disable_disconnect_push"`
 	// JoinPushMessages when enabled allow uni_websocket transport to join messages together into
 	// one frame using Centrifugal client protocol delimiters: new line for JSON protocol and
 	// length-prefixed format for Protobuf protocol. This can be useful to reduce system call

internal/timers/pool.go

@@ -0,0 +1,34 @@
+package timers
+
+import (
+	"sync"
+	"time"
+)
+
+var timerPool sync.Pool
+
+// AcquireTimer from pool.
+func AcquireTimer(d time.Duration) *time.Timer {
+	v := timerPool.Get()
+	if v == nil {
+		return time.NewTimer(d)
+	}
+	tm := v.(*time.Timer)
+	if tm.Reset(d) {
+		panic("Received an active timer from the pool!")
+	}
+	return tm
+}
+
+// ReleaseTimer to pool.
+func ReleaseTimer(tm *time.Timer) {
+	if !tm.Stop() {
+		// Collect possibly added time from the channel
+		// If timer has been stopped and nobody collected its value.
+		select {
+		case <-tm.C:
+		default:
+		}
+	}
+	timerPool.Put(tm)
+}

internal/timers/pool_test.go

@@ -0,0 +1,99 @@
+package timers
+
+import (
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// waitTimer waits for a timer to fire or a timeout, returning true if fired.
+func waitTimer(tm *time.Timer, timeout time.Duration) bool {
+	select {
+	case <-tm.C:
+		return true
+	case <-time.After(timeout):
+		return false
+	}
+}
+
+func TestTimersReliable(t *testing.T) {
+	t.Run("acquire timer does not fire immediately", func(t *testing.T) {
+		tm := AcquireTimer(50 * time.Millisecond)
+		require.NotNil(t, tm)
+
+		fired := waitTimer(tm, 10*time.Millisecond)
+		require.False(t, fired, "timer should not have fired immediately")
+
+		ReleaseTimer(tm)
+	})
+
+	t.Run("timer fires after duration", func(t *testing.T) {
+		tm := AcquireTimer(20 * time.Millisecond)
+		require.NotNil(t, tm)
+
+		fired := waitTimer(tm, 100*time.Millisecond)
+		require.True(t, fired, "timer did not fire as expected")
+
+		ReleaseTimer(tm)
+	})
+
+	t.Run("release stops timer", func(t *testing.T) {
+		tm := AcquireTimer(50 * time.Millisecond)
+		require.NotNil(t, tm)
+
+		ReleaseTimer(tm)
+
+		fired := waitTimer(tm, 60*time.Millisecond)
+		require.False(t, fired, "timer should have been stopped")
+	})
+
+	t.Run("reuse timer from pool", func(t *testing.T) {
+		// Acquire and release a timer to populate the pool
+		tm1 := AcquireTimer(50 * time.Millisecond)
+		ReleaseTimer(tm1)
+
+		// Acquire again from pool
+		tm2 := AcquireTimer(20 * time.Millisecond)
+		require.NotNil(t, tm2)
+
+		fired := waitTimer(tm2, 100*time.Millisecond)
+		require.True(t, fired, "reused timer should have fired")
+
+		ReleaseTimer(tm2)
+	})
+
+	t.Run("concurrent acquire and release", func(t *testing.T) {
+		var wg sync.WaitGroup
+		numGoroutines := 50
+
+		wg.Add(numGoroutines)
+		for i := 0; i < numGoroutines; i++ {
+			go func() {
+				defer wg.Done()
+				tm := AcquireTimer(10 * time.Millisecond)
+				// deterministically wait for timer to fire or grace period
+				waitTimer(tm, 20*time.Millisecond)
+				ReleaseTimer(tm)
+			}()
+		}
+
+		wg.Wait()
+	})
+
+	t.Run("release timer that already fired", func(t *testing.T) {
+		tm := AcquireTimer(10 * time.Millisecond)
+		require.NotNil(t, tm)
+
+		fired := waitTimer(tm, 50*time.Millisecond)
+		require.True(t, fired, "timer should have fired")
+
+		ReleaseTimer(tm)
+
+		// Acquire again to ensure pool is working
+		tm2 := AcquireTimer(20 * time.Millisecond)
+		require.NotNil(t, tm2)
+		ReleaseTimer(tm2)
+	})
+}

internal/uniws/handler.go

@@ -112,13 +112,15 @@ func (s *Handler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 	// Separate goroutine for better GC of caller's data.
 	go func() {
 		opts := websocketTransportOptions{
-			framePingInterval:  framePingInterval,
-			framePongTimeout:   framePongTimeout,
-			writeTimeout:       writeTimeout,
-			compressionMinSize: compressionMinSize,
-			pingPongConfig:     s.pingPong,
-			joinMessages:       s.config.JoinPushMessages,
-			protoMajor:         r.ProtoMajor,
+			framePingInterval:       framePingInterval,
+			framePongTimeout:        framePongTimeout,
+			writeTimeout:            writeTimeout,
+			compressionMinSize:      compressionMinSize,
+			pingPongConfig:          s.pingPong,
+			joinMessages:            s.config.JoinPushMessages,
+			protoMajor:              r.ProtoMajor,
+			disableClosingHandshake: s.config.DisableClosingHandshake,
+			disableDisconnectPush:   s.config.DisableDisconnectPush,
 		}
 
 		graceCh := make(chan struct{})
@@ -156,14 +158,45 @@ func (s *Handler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 			return nil
 		})
 
+		waitClose := func() {
+			// https://github.com/gorilla/websocket/issues/448
+			conn.SetPingHandler(nil)
+			conn.SetPongHandler(nil)
+			if s.config.DisableClosingHandshake {
+				close(graceCh)
+				return
+			}
+			_ = conn.SetReadDeadline(time.Now().Add(closeFrameWait))
+			for {
+				if _, _, err := conn.NextReader(); err != nil {
+					close(graceCh)
+					return
+				}
+			}
+		}
+
 		if req == nil {
 			_, data, err := conn.ReadMessage()
 			if err != nil {
+				waitClose()
 				return
 			}
 			_, err = json.Parse(data, &req, json.ZeroCopy)
 			if err != nil {
 				log.Info().Err(err).Str("transport", transportName).Msg("error unmarshalling connect request")
+				if !s.config.DisableClosingHandshake {
+					err = conn.WriteControl(
+						websocket.CloseMessage,
+						websocket.FormatCloseMessage(
+							int(centrifuge.DisconnectBadRequest.Code),
+							centrifuge.DisconnectBadRequest.Reason,
+						),
+						time.Now().Add(writeTimeout))
+					if err != nil {
+						return
+					}
+				}
+				waitClose()
 				return
 			}
 		}
@@ -176,16 +209,6 @@ func (s *Handler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 				break
 			}
 		}
-
-		// https://github.com/gorilla/websocket/issues/448
-		conn.SetPingHandler(nil)
-		conn.SetPongHandler(nil)
-		_ = conn.SetReadDeadline(time.Now().Add(closeFrameWait))
-		for {
-			if _, _, err := conn.NextReader(); err != nil {
-				close(graceCh)
-				break
-			}
-		}
+		waitClose()
 	}()
 }

internal/uniws/handler_test.go

@@ -3,6 +3,7 @@ package uniws
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -192,3 +193,79 @@ func TestUnidirectionalWebSocket(t *testing.T) {
 		require.True(t, pingReceived, "Expected to receive ping message")
 	})
 }
+
+func TestUnidirectionalWebSocket_CloseFrameSent(t *testing.T) {
+	t.Parallel()
+
+	// Create a custom node to trigger server-side disconnect.
+	testNode, err := centrifuge.New(centrifuge.Config{})
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = testNode.Shutdown(context.Background()) })
+
+	testNode.OnConnecting(func(ctx context.Context, event centrifuge.ConnectEvent) (centrifuge.ConnectReply, error) {
+		return centrifuge.ConnectReply{
+			Credentials: &centrifuge.Credentials{},
+		}, nil
+	})
+
+	testNode.OnConnect(func(client *centrifuge.Client) {
+		client.Disconnect(centrifuge.DisconnectConnectionLimit)
+	})
+
+	testHandler := NewHandler(testNode, Config{}, func(r *http.Request) bool {
+		return true
+	}, centrifuge.PingPongConfig{})
+
+	testServer := httptest.NewServer(middleware.LogRequest(testHandler))
+	t.Cleanup(func() { testServer.Close() })
+
+	testWsURL := "ws" + strings.TrimPrefix(testServer.URL, "http")
+
+	// Wait for test server to start.
+	for {
+		resp, err := http.Get(testServer.URL)
+		if err != nil {
+			time.Sleep(100 * time.Millisecond)
+			continue
+		}
+		_ = resp.Body.Close()
+		break
+	}
+
+	dialer := websocket.Dialer{}
+	conn, _, _, err := dialer.Dial(testWsURL, nil)
+	require.NoError(t, err)
+	defer func() { _ = conn.Close() }()
+
+	// Send connect request.
+	err = conn.WriteMessage(websocket.TextMessage, []byte(`{}`))
+	require.NoError(t, err)
+
+	// Read connect reply.
+	_, data, err := conn.ReadMessage()
+	require.NoError(t, err)
+	ensureMessageHasClient(t, data)
+
+	require.NoError(t, conn.SetReadDeadline(time.Now().Add(10*time.Second)))
+
+	// Read until connection is closed - server should send close frame.
+	var closeErr *websocket.CloseError
+	for {
+		_, _, err := conn.ReadMessage()
+		if err != nil {
+			t.Logf("ReadMessage error: %v", err)
+			closeErr = &websocket.CloseError{}
+			if ok := errors.As(err, &closeErr); ok {
+				t.Logf("Close frame received: code=%d, text=%s", closeErr.Code, closeErr.Text)
+				break
+			}
+			// Some other error occurred.
+			break
+		}
+	}
+
+	// Verify close frame was received with correct code and reason.
+	require.NotNil(t, closeErr, "Expected to receive close frame")
+	require.Equal(t, centrifuge.DisconnectConnectionLimit.Code, uint32(closeErr.Code))
+	require.Equal(t, centrifuge.DisconnectConnectionLimit.Reason, closeErr.Text)
+}