267 lines
7.0 KiB
Go
267 lines
7.0 KiB
Go
package uniws
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/centrifugal/centrifugo/v6/internal/timers"
|
|
"github.com/centrifugal/centrifugo/v6/internal/tools"
|
|
"github.com/centrifugal/centrifugo/v6/internal/websocket"
|
|
|
|
"github.com/centrifugal/centrifuge"
|
|
"github.com/centrifugal/protocol"
|
|
)
|
|
|
|
// websocketTransport is a wrapper struct over websocket connection to fit session
|
|
// interface so client will accept it.
|
|
type websocketTransport struct {
|
|
mu sync.RWMutex
|
|
writeMu sync.Mutex // sync general write with unidirectional ping write.
|
|
conn *websocket.Conn
|
|
closeCh chan struct{}
|
|
graceCh chan struct{}
|
|
opts websocketTransportOptions
|
|
pingTimer *time.Timer
|
|
closed bool
|
|
}
|
|
|
|
type websocketTransportOptions struct {
|
|
framePingInterval time.Duration
|
|
framePongTimeout time.Duration
|
|
writeTimeout time.Duration
|
|
compressionMinSize int
|
|
pingPongConfig centrifuge.PingPongConfig
|
|
protoMajor int
|
|
joinMessages bool
|
|
disableClosingHandshake bool // TODO: consider keeping/removing for Centrifugo v7.
|
|
disableDisconnectPush bool
|
|
}
|
|
|
|
func newWebsocketTransport(conn *websocket.Conn, opts websocketTransportOptions, graceCh chan struct{}) *websocketTransport {
|
|
transport := &websocketTransport{
|
|
conn: conn,
|
|
closeCh: make(chan struct{}),
|
|
graceCh: graceCh,
|
|
opts: opts,
|
|
}
|
|
if opts.framePingInterval > 0 {
|
|
transport.addPing()
|
|
}
|
|
return transport
|
|
}
|
|
|
|
func (t *websocketTransport) ping() {
|
|
select {
|
|
case <-t.closeCh:
|
|
return
|
|
default:
|
|
if t.opts.framePongTimeout > 0 {
|
|
// It's safe to call SetReadDeadline concurrently with reader in separate goroutine.
|
|
// According to Go docs:
|
|
// SetReadDeadline sets the deadline for future Read calls and any currently-blocked Read call.
|
|
_ = t.conn.SetReadDeadline(time.Now().Add(t.opts.framePongTimeout))
|
|
}
|
|
err := t.conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(t.opts.writeTimeout))
|
|
if err != nil {
|
|
_ = t.Close(centrifuge.DisconnectWriteError)
|
|
return
|
|
}
|
|
if t.opts.writeTimeout > 0 && t.opts.protoMajor > 1 {
|
|
// WriteControl set a write deadline on the underlying HTTP/2 stream but,
|
|
// unlike writeData, nothing clears it here. An expired HTTP/2 deadline
|
|
// cannot be extended and permanently RSTs the stream, so an idle
|
|
// connection would be dropped ~writeTimeout after each ping. Clear it,
|
|
// serialized against writeData via writeMu.
|
|
t.writeMu.Lock()
|
|
derr := t.conn.NetConn().SetWriteDeadline(time.Time{})
|
|
t.writeMu.Unlock()
|
|
if derr != nil {
|
|
_ = t.Close(centrifuge.DisconnectWriteError)
|
|
return
|
|
}
|
|
}
|
|
t.addPing()
|
|
}
|
|
}
|
|
|
|
func (t *websocketTransport) addPing() {
|
|
t.mu.Lock()
|
|
if t.closed {
|
|
t.mu.Unlock()
|
|
return
|
|
}
|
|
t.pingTimer = time.AfterFunc(t.opts.framePingInterval, t.ping)
|
|
t.mu.Unlock()
|
|
}
|
|
|
|
const transportName = "uni_websocket"
|
|
|
|
// Name returns name of transport.
|
|
func (t *websocketTransport) Name() string {
|
|
return transportName
|
|
}
|
|
|
|
// Protocol returns transport protocol.
|
|
func (t *websocketTransport) Protocol() centrifuge.ProtocolType {
|
|
return centrifuge.ProtocolTypeJSON
|
|
}
|
|
|
|
// ProtocolVersion returns transport protocol version.
|
|
func (t *websocketTransport) ProtocolVersion() centrifuge.ProtocolVersion {
|
|
return centrifuge.ProtocolVersion2
|
|
}
|
|
|
|
// Unidirectional returns whether transport is unidirectional.
|
|
func (t *websocketTransport) Unidirectional() bool {
|
|
return true
|
|
}
|
|
|
|
// DisabledPushFlags ...
|
|
func (t *websocketTransport) DisabledPushFlags() uint64 {
|
|
if t.opts.disableDisconnectPush {
|
|
return centrifuge.PushFlagDisconnect
|
|
}
|
|
return 0
|
|
}
|
|
|
|
// PingPongConfig ...
|
|
func (t *websocketTransport) PingPongConfig() centrifuge.PingPongConfig {
|
|
return t.opts.pingPongConfig
|
|
}
|
|
|
|
// Emulation ...
|
|
func (t *websocketTransport) Emulation() bool {
|
|
return false
|
|
}
|
|
|
|
// AcceptProtocol ...
|
|
func (t *websocketTransport) AcceptProtocol() string {
|
|
return tools.GetAcceptProtocolLabel(t.opts.protoMajor)
|
|
}
|
|
|
|
func (t *websocketTransport) writeData(data []byte) error {
|
|
if t.opts.compressionMinSize > 0 {
|
|
t.conn.EnableWriteCompression(len(data) > t.opts.compressionMinSize)
|
|
}
|
|
var messageType = websocket.TextMessage
|
|
if t.Protocol() == centrifuge.ProtocolTypeProtobuf {
|
|
messageType = websocket.BinaryMessage
|
|
}
|
|
|
|
t.writeMu.Lock()
|
|
if t.opts.writeTimeout > 0 {
|
|
_ = t.conn.SetWriteDeadline(time.Now().Add(t.opts.writeTimeout))
|
|
}
|
|
err := t.conn.WriteMessage(messageType, data)
|
|
if err != nil {
|
|
t.writeMu.Unlock()
|
|
return err
|
|
}
|
|
if t.opts.writeTimeout > 0 {
|
|
_ = t.conn.SetWriteDeadline(time.Time{})
|
|
if t.opts.protoMajor > 1 {
|
|
// For HTTP/2 connections, we need to actually clear the deadline on the underlying
|
|
// connection. The websocket Conn.SetWriteDeadline only sets a field, but doesn't
|
|
// clear the deadline that was already set on the underlying net.Conn during write.
|
|
// This is critical for HTTP/2 ResponseController where expired deadlines cannot be
|
|
// extended and will cause the stream to fail permanently.
|
|
err = t.conn.NetConn().SetWriteDeadline(time.Time{})
|
|
if err != nil {
|
|
t.writeMu.Unlock()
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
t.writeMu.Unlock()
|
|
|
|
return nil
|
|
}
|
|
|
|
func (t *websocketTransport) writeMany(messages ...[]byte) error {
|
|
protoType := protocol.TypeJSON
|
|
if t.Protocol() == centrifuge.ProtocolTypeProtobuf {
|
|
protoType = protocol.TypeProtobuf
|
|
}
|
|
encoder := protocol.GetDataEncoder(protoType)
|
|
defer protocol.PutDataEncoder(protoType, encoder)
|
|
for i := range messages {
|
|
_ = encoder.Encode(messages[i])
|
|
}
|
|
return t.writeData(encoder.Finish())
|
|
}
|
|
|
|
// Write data to transport.
|
|
func (t *websocketTransport) Write(message []byte) error {
|
|
return t.WriteMany(message)
|
|
}
|
|
|
|
// WriteMany data to transport.
|
|
func (t *websocketTransport) WriteMany(messages ...[]byte) error {
|
|
select {
|
|
case <-t.closeCh:
|
|
return nil
|
|
default:
|
|
var err error
|
|
if t.opts.joinMessages {
|
|
if len(messages) == 1 && t.Protocol() == centrifuge.ProtocolTypeJSON {
|
|
// Fast path for one JSON message.
|
|
return t.writeData(messages[0])
|
|
}
|
|
err = t.writeMany(messages...)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
} else {
|
|
for i := 0; i < len(messages); i++ {
|
|
err = t.writeData(messages[i])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
}
|
|
|
|
const closeFrameWait = 5 * time.Second
|
|
|
|
// Close closes transport.
|
|
func (t *websocketTransport) Close(disconnect centrifuge.Disconnect) error {
|
|
t.mu.Lock()
|
|
if t.closed {
|
|
t.mu.Unlock()
|
|
return nil
|
|
}
|
|
t.closed = true
|
|
if t.pingTimer != nil {
|
|
t.pingTimer.Stop()
|
|
}
|
|
close(t.closeCh)
|
|
t.mu.Unlock()
|
|
|
|
if t.opts.disableClosingHandshake {
|
|
return t.conn.Close()
|
|
}
|
|
|
|
if disconnect.Code != centrifuge.DisconnectConnectionClosed.Code {
|
|
msg := websocket.FormatCloseMessage(int(disconnect.Code), disconnect.Reason)
|
|
err := t.conn.WriteControl(websocket.CloseMessage, msg, time.Now().Add(t.opts.writeTimeout))
|
|
if err != nil {
|
|
return t.conn.Close()
|
|
}
|
|
select {
|
|
case <-t.graceCh:
|
|
default:
|
|
// Wait for closing handshake completion.
|
|
tm := timers.AcquireTimer(closeFrameWait)
|
|
select {
|
|
case <-t.graceCh:
|
|
case <-tm.C:
|
|
}
|
|
timers.ReleaseTimer(tm)
|
|
}
|
|
return t.conn.Close()
|
|
}
|
|
return t.conn.Close()
|
|
}
|