connect/server/client.go

package server

import (
	"encoding/json"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"sikey.com/websocket/models"
	"sikey.com/websocket/repositories"
	"sikey.com/websocket/utils/zlog"
)

type Client struct {
	ctx            *gin.Context
	UserId         string
	hub            *Hub
	UnderlyingConn *websocket.Conn

	isSimpleMsg  bool   // isSimpleMsg 是否是简单消息
	localization string // localization 国际码

	// Send message channel 发送消息
	// 当用户在线时会通过 Send channel 发送在线消息 但如果用户不在线,
	// 则会通过 SendOffline channel 发送离线消息
	Send        chan *Message
	SendOffline chan *Message // SendOffline 发送离线消息

	readWait  time.Duration // readWait 读超时
	writeWait time.Duration // writeWait 写超时

	pingWait time.Duration
	pongWait time.Duration

	repos *repositories.Repositories
}

func (c *Client) reader() {
	defer func() {
		c.hub.Disconnect <- c
		c.UnderlyingConn.Close()
	}()

	c.UnderlyingConn.SetReadDeadline(time.Now().Add(c.readWait))

	for {

		_, bytes, err := c.UnderlyingConn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				zlog.Errorf("error: %v", err)
			}
			break
		}

		message := deserializeMessage(bytes)
		switch message.Type {
		case MessageTypePing:

			// Reset read dead line, prevent Reader from shutting down
			c.UnderlyingConn.SetReadDeadline(time.Now().Add(c.pongWait))

			// Reply pong message
			c.Send <- pongMessage()

		case MessageTypeChating:

			// Chat dialogue messages
			chatingContent := message.Content.(ChatingContent)

			// Save message to database
			payload, _ := json.Marshal(chatingContent.Payload)
			err := c.repos.SessionRepository.CreateMessage(c.ctx, &models.SessionMessage{
				ID:          message.MessageId,
				SessionId:   chatingContent.Receiver,
				Receiver:    chatingContent.Receiver,
				Sender:      c.UserId,
				Type:        uint8(message.Type),
				ContentType: uint8(chatingContent.PayloadType),
				Content:     string(payload),
				IsRead:      false,
			})

			if err != nil {
				c.writeError(message.MessageId, err)
				continue
			}

			members, err := c.repos.SessionRepository.GetSessionMembersRemoveOneself(
				c.ctx, chatingContent.Receiver, c.UserId)
			if err != nil {
				c.writeError(message.MessageId, err)
				continue
			}

			for _, member := range members {
				var messaging = *message
				messaging.receiver = member.AccountId
				messaging.Content = chatingContent
				c.hub.Message <- &messaging
			}
		}

		// Reply message
		if message.IsNeedReply() {
			c.Send <- newReplyMessage(message.MessageId)
		}
	}
}

func (c *Client) writer() {
	ticker := time.NewTicker(c.pingWait)
	defer func() {
		ticker.Stop()
		c.UnderlyingConn.Close()
	}()

	for {
		select {
		case message, ok := <-c.Send:
			c.UnderlyingConn.SetWriteDeadline(time.Now().Add(c.writeWait))
			if !ok {
				// The hub closed the channel.
				c.UnderlyingConn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			var err error
			if message.Type == MessageTypePong {
				err = c.UnderlyingConn.WriteMessage(websocket.TextMessage, []byte("pong"))
				if err != nil {
					return
				}
				break
			}

			err = c.UnderlyingConn.WriteMessage(websocket.TextMessage, serializationMessage(message))
			if err != nil {
				return
			}
		case <-ticker.C:
			// 到时间发送 ping 信号
			c.UnderlyingConn.SetWriteDeadline(time.Now().Add(c.writeWait))
			if err := c.UnderlyingConn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

func (c *Client) writeError(messageId string, err error) {
	c.Send <- &Message{
		MessageId: messageId,
		Type:      MessageTypeError,
		Content:   ContentError{Err: err.Error()},
	}
}