connect/server/server.go

package server

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

	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt/v5"
	"github.com/gorilla/websocket"
	"github.com/rotisserie/eris"
	"github.com/spf13/viper"
	"go.uber.org/zap"
	"sikey.com/websocket/config"
	"sikey.com/websocket/models"
	"sikey.com/websocket/repositories"
	"sikey.com/websocket/utils/keys"
)

type Server struct {
	ID           string
	Ctx          *gin.Context
	Repositories *repositories.Repositories

	Upgrader websocket.Upgrader
	Hub      *Hub
}

func (srv *Server) WebsocketHandler(ctx *gin.Context) {
	srv.Ctx = ctx

	// Builder headers
	headers := headerBuilder(ctx)

	// Validate token
	id, ok, err := jwtParse(headers)
	if !ok {
		zap.L().Error("invalid token", zap.Error(err))
		ctx.AbortWithError(http.StatusUnauthorized, err)
		return
	}

	zap.L().Info("client online", zap.String("user_id", id))
	conn, err := srv.Upgrader.Upgrade(ctx.Writer, ctx.Request, nil)
	if err != nil {
		zap.L().Error("upgrade error", zap.Error(err))
		ctx.AbortWithError(http.StatusInternalServerError, err)
		return
	}

	// Reconnection mechanism
	var client *Client
	if client = srv.Hub.getClientByUserId(id); client != nil {
		_ = client.UnderlyingConn.Close()
		client.UnderlyingConn = conn
		zap.L().Info("client reconnection", zap.String("user_id", id))
		return
	}

	// Create client
	client = &Client{
		ctx:            ctx.Copy(),
		UserId:         id,
		hub:            srv.Hub,
		UnderlyingConn: conn,
		online:         &models.Online{UserId: id, ServerId: srv.ID},
		Send:           make(chan *Message, config.Websocket.MessageSize),
		writeWait:      config.Websocket.WriteWait * time.Second,
		readWait:       config.Websocket.ReadWait * time.Second,
		pingWait:       config.Websocket.HeartbeatWait * time.Second,

		isSimpleMsg:  headers[keys.SimpleHeader].(bool),
		localization: headers[keys.LocalizationHeader].(string),

		repos: srv.Repositories,
	}
	srv.Hub.Connect <- client

	// Online status to redis
	online := &models.Online{UserId: client.UserId, ServerId: srv.ID}
	if err := srv.Repositories.OnlineRepository.SetOnline(ctx, online); err != nil {
		ctx.AbortWithError(http.StatusInternalServerError,
			eris.Wrapf(err, "unable to set online status for user: %s", client.UserId))
		return
	}

	go client.reader()
	go client.writer()
}

type Headers = map[string]interface{}

func headerBuilder(ctx *gin.Context) Headers {
	headers := make(Headers)
	request := ctx.Request
	accessToken := request.URL.Query().Get(keys.AccessTokenHeader)
	simple := request.URL.Query().Get(keys.SimpleHeader)
	localization := request.URL.Query().Get(keys.LocalizationHeader)
	headers[keys.UserIdHeader] = request.URL.Query().Get(keys.UserIdHeader)
	headers[keys.AccessTokenHeader] = accessToken
	headers[keys.SimpleHeader] = simple == "1"
	headers[keys.LocalizationHeader] = localization
	return headers
}

type UserClaims struct {
	jwt.RegisteredClaims

	UID string `json:"uid"`
}

func jwtParse(headers Headers) (string, bool, error) {
	if userId, ok := headers[keys.UserIdHeader]; ok {
		if userId != "" {
			return userId.(string), true, nil
		}
	}

	accessToken := headers[keys.AccessTokenHeader].(string)
	if len(accessToken) == 0 {
		return "", false, eris.New("token is empty")
	}
	accessToken = strings.Trim(accessToken, " ")

	token, err := jwt.ParseWithClaims(accessToken, &UserClaims{}, func(t *jwt.Token) (interface{}, error) {
		return []byte(viper.GetString("auth.secret")), nil
	})
	userClaims, ok := token.Claims.(*UserClaims)
	if ok {
		return userClaims.UID, true, nil
	} else {
		return "", false, eris.Wrap(err, "token parse error")
	}
	// mapClaims := token.Claims.(jwt.MapClaims)
	// exp := mapClaims["exp"].(float64)
	// if exp != 0 && exp < float64(time.Now().Unix()) {
	// 	return "", false, eris.New("token is expired")
	// }
}