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"
	"sikey.com/websocket/config"
	"sikey.com/websocket/models"
	"sikey.com/websocket/repositories"
	"sikey.com/websocket/utils/keys"
	"sikey.com/websocket/utils/zlog"
)

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

	conn, err := srv.Upgrader.Upgrade(ctx.Writer, ctx.Request, nil)
	if err != nil {
		zlog.Error(err)
		ctx.AbortWithError(http.StatusInternalServerError, err)
		return
	}

	// Builder headers
	headers := headerBuilder(ctx)

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

	// Create client
	client := &Client{
		ctx:            ctx.Copy(),
		UserId:         id,
		hub:            srv.Hub,
		UnderlyingConn: conn,
		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
	zlog.Debugf("client: %s", client.UserId)

	// 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
}

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.Parse(accessToken, func(token *jwt.Token) (interface{}, error) {
		return []byte(viper.GetString("auth.secret")), nil
	})
	if err != nil {
		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")
	// }
	uid := mapClaims["Uid"].(string)
	return uid, true, nil
}