im_server/server/hdl_websock.go

/******************************************************************************
 *
 *  Description :
 *
 *    Handler of websocket connections. See also hdl_longpoll.go for long polling
 *    and hdl_grpc.go for gRPC.
 *
 *****************************************************************************/

package main

import (
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/gorilla/websocket"
	"github.com/tinode/chat/server/logs"
	immetrics "github.com/tinode/chat/server/metrics"
	"github.com/tinode/chat/server/store/types"
)

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second
	// Heartbeat defaults tolerate transient mobile and public-network stalls while
	// still detecting dead connections in a bounded amount of time.
	defaultWSPongWait   = 90 * time.Second
	defaultWSPingPeriod = 15 * time.Second
)

var (
	pongWait   = defaultWSPongWait
	pingPeriod = defaultWSPingPeriod
)

func configureWSHeartbeat(pingSeconds, pongSeconds int) error {
	ping := defaultWSPingPeriod
	pong := defaultWSPongWait
	if pingSeconds > 0 {
		ping = time.Duration(pingSeconds) * time.Second
	}
	if pongSeconds > 0 {
		pong = time.Duration(pongSeconds) * time.Second
	}
	if ping >= pong {
		return errors.New("websocket ping period must be shorter than pong wait")
	}
	pingPeriod = ping
	pongWait = pong
	immetrics.SetWSHeartbeatConfig(ping, pong)
	return nil
}

func (sess *Session) closeWS() {
	if sess.proto == WEBSOCK {
		sess.ws.Close()
	}
}

func (sess *Session) readLoop() {
	defer func() {
		sess.closeWS()
		sess.cleanUp(false)
	}()

	sess.ws.SetReadLimit(globals.maxMessageSize)
	refreshReadDeadline := func(event, presenceSource string) {
		atomic.StoreInt64(&sess.lastAction, types.TimeNow().UnixNano())
		immetrics.IncWSHeartbeat(event)
		if presenceSource != "" {
			sess.refreshPresence(presenceSource)
		}
		_ = sess.ws.SetReadDeadline(time.Now().Add(pongWait))
	}
	_ = sess.ws.SetReadDeadline(time.Now().Add(pongWait))
	sess.ws.SetPongHandler(func(string) error {
		refreshReadDeadline("pong_received", "ws.pong")
		return nil
	})
	sess.ws.SetPingHandler(func(appData string) error {
		refreshReadDeadline("ping_received", "")
		return sess.ws.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(writeWait))
	})

	for {
		// Read a ClientComMessage
		_, raw, err := sess.ws.ReadMessage()
		if err != nil {
			reason := classifyWSError(err)
			immetrics.IncWSDisconnect(reason)
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				logs.Err.Printf("ws: readLoop sid=%s uid=%s reason=%s err=%v", sess.sid, sess.uid.UserId(), reason, err)
			}
			return
		}
		// Any successfully read client frame proves that the connection is alive.
		// This prevents a busy client from timing out merely because its Pong was
		// delayed behind application traffic.
		atomic.StoreInt64(&sess.lastAction, types.TimeNow().UnixNano())
		_ = sess.ws.SetReadDeadline(time.Now().Add(pongWait))
		statsInc("IncomingMessagesWebsockTotal", 1)
		dispatchStarted := time.Now()
		sess.dispatchRaw(raw)
		immetrics.ObserveWSReadDispatch(dispatchStarted)
	}
}

func (sess *Session) sendMessage(msg any) bool {
	if depth := len(sess.send); depth > sendQueueLimit {
		immetrics.ObserveWSWrite("text", "rejected", "queue_limit", time.Now(), 0)
		logs.Err.Printf("ws: outbound queue limit exceeded sid=%s uid=%s depth=%d capacity=%d",
			sess.sid, sess.uid.UserId(), depth, cap(sess.send))
		return false
	}

	statsInc("OutgoingMessagesWebsockTotal", 1)
	if err := sess.writeWSFrame(websocket.TextMessage, msg); err != nil {
		return false
	}
	return true
}

func (sess *Session) writeLoop() {
	ticker := time.NewTicker(pingPeriod)

	defer func() {
		ticker.Stop()
		// Break readLoop.
		sess.closeWS()
	}()

	for {
		select {
		case msg, ok := <-sess.send:
			if !ok {
				// Channel closed.
				return
			}
			sess.observeOutboundDequeue(false)
			switch v := msg.(type) {
			case []*ServerComMessage: // batch of unserialized messages
				for _, msg := range v {
					w := sess.serializeAndUpdateStats(msg)
					if !sess.sendMessage(w) {
						return
					}
					if isAcceptedCtrl(msg) {
						immetrics.ObserveCtrlACKDelivery("write", msg.Timestamp)
					}
				}
			case *ServerComMessage: // single unserialized message
				w := sess.serializeAndUpdateStats(v)
				if !sess.sendMessage(w) {
					return
				}
				if isAcceptedCtrl(v) {
					immetrics.ObserveCtrlACKDelivery("write", v.Timestamp)
				}
			case sharedServerMessage:
				statsAddHistSample("OutgoingMessageSize", float64(len(v)))
				if !sess.sendMessage([]byte(v)) {
					return
				}
			default: // serialized message
				if !sess.sendMessage(v) {
					return
				}
			}

		case <-sess.bkgTimer.C:
			if sess.background {
				sess.background = false
				sess.onBackgroundTimer()
			}

		case msg := <-sess.stop:
			// Shutdown requested, don't care if the message is delivered
			if msg != nil {
				sess.writeWSFrame(websocket.TextMessage, msg)
			}
			return

		case topic := <-sess.detach:
			sess.delSub(topic)

		case <-ticker.C:
			if err := sess.writeWSFrame(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

func (sess *Session) writeWSFrame(mt int, msg any) error {
	frame := "text"
	if mt == websocket.PingMessage {
		frame = "ping"
	}
	payloadBytes := 0
	if bits, ok := msg.([]byte); ok {
		payloadBytes = len(bits)
	}
	started := time.Now()
	err := wsWrite(sess.ws, mt, msg)
	if err == nil {
		immetrics.ObserveWSWrite(frame, "ok", "none", started, payloadBytes)
		return nil
	}
	reason := classifyWSError(err)
	immetrics.ObserveWSWrite(frame, "error", reason, started, payloadBytes)
	if shouldLogWSWriteError(reason) {
		logs.Err.Printf("ws: writeLoop sid=%s uid=%s frame=%s queue=%d capacity=%d reason=%s err=%v",
			sess.sid, sess.uid.UserId(), frame, len(sess.send), cap(sess.send), reason, err)
	}
	return err
}

func shouldLogWSWriteError(reason string) bool {
	switch reason {
	case "close_normal", "close_going_away", "closed", "eof", "broken_pipe", "connection_reset":
		return false
	default:
		return true
	}
}

func classifyWSError(err error) string {
	if err == nil {
		return "none"
	}
	var closeErr *websocket.CloseError
	if errors.As(err, &closeErr) {
		switch closeErr.Code {
		case websocket.CloseNormalClosure:
			return "close_normal"
		case websocket.CloseGoingAway:
			return "close_going_away"
		case websocket.CloseAbnormalClosure:
			return "close_abnormal"
		default:
			return "close_other"
		}
	}
	var netErr net.Error
	if errors.As(err, &netErr) && netErr.Timeout() {
		return "timeout"
	}
	switch {
	case errors.Is(err, syscall.ECONNRESET):
		return "connection_reset"
	case errors.Is(err, syscall.EPIPE):
		return "broken_pipe"
	case errors.Is(err, net.ErrClosed):
		return "closed"
	case errors.Is(err, io.EOF):
		return "eof"
	default:
		return "other"
	}
}

// Writes a message with the given message type (mt) and payload.
func wsWrite(ws *websocket.Conn, mt int, msg any) error {
	var bits []byte
	if msg != nil {
		bits = msg.([]byte)
	} else {
		bits = []byte{}
	}
	ws.SetWriteDeadline(time.Now().Add(writeWait))
	return ws.WriteMessage(mt, bits)
}

// Handles websocket requests from peers.
var upgrader = websocket.Upgrader{
	ReadBufferSize:    1024,
	WriteBufferSize:   1024,
	EnableCompression: globals.wsCompression,
	// Allow connections from any Origin
	CheckOrigin: func(r *http.Request) bool { return true },
}

func serveWebSocket(wrt http.ResponseWriter, req *http.Request) {
	now := types.TimeNow()

	if isValid, _ := checkAPIKey(getAPIKey(req)); !isValid {
		wrt.WriteHeader(http.StatusForbidden)
		json.NewEncoder(wrt).Encode(ErrAPIKeyRequired(now))
		logs.Err.Println("ws: Missing, invalid or expired API key")
		return
	}

	if req.Method != http.MethodGet {
		wrt.WriteHeader(http.StatusMethodNotAllowed)
		json.NewEncoder(wrt).Encode(ErrOperationNotAllowed("", "", now))
		logs.Err.Println("ws: Invalid HTTP method", req.Method)
		return
	}

	ws, err := upgrader.Upgrade(wrt, req, nil)
	if _, ok := err.(websocket.HandshakeError); ok {
		logs.Err.Println("ws: Not a websocket handshake")
		return
	} else if err != nil {
		logs.Err.Println("ws: failed to Upgrade ", err)
		return
	}

	sess, count := globals.sessionStore.NewSession(ws, "")
	if globals.useXForwardedFor {
		sess.remoteAddr = req.Header.Get("X-Forwarded-For")
		if !isRoutableIP(sess.remoteAddr) {
			sess.remoteAddr = ""
		}
	}
	if sess.remoteAddr == "" {
		sess.remoteAddr = req.RemoteAddr
	}

	logs.Info.Println("ws: session started", sess.sid, sess.remoteAddr, count)

	// Do work in goroutines to return from serveWebSocket() to release file pointers.
	// Otherwise "too many open files" will happen.
	go sess.writeLoop()
	go sess.readLoop()
}