screego/ws/client.go

package ws

import (
	"fmt"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/gorilla/websocket"
	"github.com/rs/xid"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
	"github.com/screego/server/ws/outgoing"
)

var ping = func(conn *websocket.Conn) error {
	return conn.WriteMessage(websocket.PingMessage, nil)
}

var writeJSON = func(conn *websocket.Conn, v interface{}) error {
	return conn.WriteJSON(v)
}

const (
	writeWait = 2 * time.Second
)

type Client struct {
	conn *websocket.Conn
	info ClientInfo
	once once
	read chan<- ClientMessage
}

type ClientMessage struct {
	Info               ClientInfo
	SkipConnectedCheck bool
	Incoming           Event
}

type ClientInfo struct {
	ID                xid.ID
	Authenticated     bool
	AuthenticatedUser string
	Write             chan outgoing.Message
	Addr              net.IP
}

func newClient(conn *websocket.Conn, req *http.Request, read chan ClientMessage, authenticatedUser string, authenticated, trustProxy bool) *Client {
	ip := conn.RemoteAddr().(*net.TCPAddr).IP
	if realIP := req.Header.Get("X-Real-IP"); trustProxy && realIP != "" {
		ip = net.ParseIP(realIP)
	}

	client := &Client{
		conn: conn,
		info: ClientInfo{
			Authenticated:     authenticated,
			AuthenticatedUser: authenticatedUser,
			ID:                xid.New(),
			Addr:              ip,
			Write:             make(chan outgoing.Message, 1),
		},
		read: read,
	}
	client.debug().Msg("WebSocket New Connection")
	return client
}

// CloseOnError closes the connection.
func (c *Client) CloseOnError(code int, reason string) {
	c.once.Do(func() {
		go func() {
			c.read <- ClientMessage{
				Info: c.info,
				Incoming: &Disconnected{
					Code:   code,
					Reason: reason,
				},
			}
		}()
		c.writeCloseMessage(code, reason)
	})
}

func (c *Client) CloseOnDone(code int, reason string) {
	c.once.Do(func() {
		c.writeCloseMessage(code, reason)
	})
}

func (c *Client) writeCloseMessage(code int, reason string) {
	message := websocket.FormatCloseMessage(code, reason)
	_ = c.conn.WriteControl(websocket.CloseMessage, message, time.Now().Add(writeWait))
	c.conn.Close()
}

// startWriteHandler starts listening on the client connection. As we do not need anything from the client,
// we ignore incoming messages. Leaves the loop on errors.
func (c *Client) startReading(pongWait time.Duration) {
	defer c.CloseOnError(websocket.CloseNormalClosure, "Reader Routine Closed")

	_ = c.conn.SetReadDeadline(time.Now().Add(pongWait))
	c.conn.SetPongHandler(func(appData string) error {
		_ = c.conn.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})
	for {
		t, m, err := c.conn.NextReader()
		if err != nil {
			c.CloseOnError(websocket.CloseNormalClosure, "read error: "+err.Error())
			return
		}
		if t == websocket.BinaryMessage {
			c.CloseOnError(websocket.CloseUnsupportedData, "unsupported binary message type")
			return
		}

		incoming, err := ReadTypedIncoming(m)
		if err != nil {
			c.CloseOnError(websocket.CloseUnsupportedData, fmt.Sprintf("malformed message: %s", err))
			return
		}
		c.debug().Interface("event", fmt.Sprintf("%T", incoming)).Interface("payload", incoming).Msg("WebSocket Receive")
		c.read <- ClientMessage{Info: c.info, Incoming: incoming}
	}
}

// startWriteHandler starts the write loop. The method has the following tasks:
// * ping the client in the interval provided as parameter
// * write messages send by the channel to the client
// * on errors exit the loop.
func (c *Client) startWriteHandler(pingPeriod time.Duration) {
	pingTicker := time.NewTicker(pingPeriod)
	defer pingTicker.Stop()
	defer func() {
		c.debug().Msg("WebSocket Done")
	}()
	defer c.conn.Close()
	for {
		select {
		case message := <-c.info.Write:
			if msg, ok := message.(outgoing.CloseWriter); ok {
				c.debug().Str("reason", msg.Reason).Int("code", msg.Code).Msg("WebSocket Close")
				c.CloseOnDone(msg.Code, msg.Reason)
				return
			}

			_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			typed, err := ToTypedOutgoing(message)
			c.debug().Interface("event", typed.Type).Interface("payload", typed.Payload).Msg("WebSocket Send")
			if err != nil {
				c.debug().Err(err).Msg("could not get typed message, exiting connection.")
				c.CloseOnError(websocket.CloseNormalClosure, "malformed outgoing "+err.Error())
				continue
			}

			if err := writeJSON(c.conn, typed); err != nil {
				c.printWebSocketError("write", err)
				c.CloseOnError(websocket.CloseNormalClosure, "write error"+err.Error())
			}
		case <-pingTicker.C:
			_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := ping(c.conn); err != nil {
				c.printWebSocketError("ping", err)
				c.CloseOnError(websocket.CloseNormalClosure, "ping timeout")
			}
		}
	}
}

func (c *Client) debug() *zerolog.Event {
	return log.Debug().Str("id", c.info.ID.String()).Str("ip", c.info.Addr.String())
}

func (c *Client) printWebSocketError(typex string, err error) {
	if strings.Contains(err.Error(), "use of closed network connection") {
		return
	}
	closeError, ok := err.(*websocket.CloseError)

	if ok && closeError != nil && (closeError.Code == 1000 || closeError.Code == 1001) {
		// normal closure
		return
	}

	c.debug().Str("type", typex).Err(err).Msg("WebSocket Error")
}