nn.daxia.dev/router/ws_server.go

package router

import (
	"encoding/base64"
	"errors"
	"log"
	"net/http"
	"runtime"

	jsoniter "github.com/json-iterator/go"
	"github.com/sirupsen/logrus"
	"gogs.daxia.dev/huanan/pkg.daxia.dev/encrypt"
	"nn.daxia.dev/game/qznn"
	"nn.daxia.dev/handler/api"
	"nn.daxia.dev/model"

	"github.com/gorilla/websocket"
)

//TODO:限制重复链接的问题

type BaseReq struct {
	Type string `json:"type"`
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 5096,

	// 解决跨域问题
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func Index(w http.ResponseWriter, r *http.Request) {
	defer func() {
		if err := recover(); err != nil {
			logrus.Error(err)
			buf := make([]byte, 2048*5)
			n := runtime.Stack(buf, false)
			stackInfo := string(buf[:n])
			logrus.Errorf("panic stack info %s", stackInfo)
		}
	}()

	var err error

	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}

	defer c.Close()

	debugAuth := r.URL.Query().Get("debug_auth")
	aesKey := model.AESKey
	if debugAuth == model.DebugStr {
		aesKey = nil
	}

	var ctx map[string]interface{}

	for {
		err = func() error {
			defer func() {
				if err := recover(); err != nil {
					logrus.Error(err)
					buf := make([]byte, 2048*5)
					n := runtime.Stack(buf, false)
					stackInfo := string(buf[:n])
					logrus.Errorf("panic stack info %s", stackInfo)
				}
			}()

			_, message, err := c.ReadMessage()
			if err != nil {
				logrus.Error("read:", err)
				return err
			}

			baseReq, message, err := parseMsg(message, aesKey)
			if err != nil {
				logrus.Error(err)
				return err
			}

			if baseReq.Type != "ping" {
				logrus.Infof("new msg income:%s", baseReq.Type)
			}

			switch baseReq.Type {
			case "perform":
				err = api.Perform(c, message, ctx)
			case "ping":
				err = api.GamePingMessage(c, message, ctx)
			case "connect":
				ctx, err = GameAuthMiddleware(c, message, api.GameConnect)
			case "disconnect":
				api.GameDisconnect(c, nil, ctx)
				return errors.New("disconnect")
			}

			if err != nil {
				logrus.Error(err, " 消息:", string(message))
				return nil
			}

			userModelI, exists := ctx["userModel"]
			if exists {
				userModel := userModelI.(model.User)
				if !qznn.UserIsOnline(userModel.ID) {
					return errors.New("用户掉线")
				}
			}

			return nil
		}()

		if err != nil {
			logrus.Infof("close: %v, err:%s", ctx, err.Error())
			break
		}
	}

	api.GameDisconnect(c, nil, ctx)
}

func parseMsg(message, aesKey []byte) (*BaseReq, []byte, error) {
	decBodyStr := string(message)

	if aesKey != nil {
		decBody, err := base64.StdEncoding.DecodeString(string(decBodyStr))
		if err != nil {
			logrus.Error(err, ":", string(decBodyStr))
			return nil, nil, err
		}

		decBody, err = encrypt.AesDecrypt(decBody, aesKey)
		if err != nil {
			logrus.Error(err)
			return nil, nil, err
		}

		decBodyStr = string(decBody)
	}

	baseReq := BaseReq{}
	err := jsoniter.Unmarshal([]byte(decBodyStr), &baseReq)
	if err != nil {
		return nil, nil, err
	}

	return &baseReq, []byte(decBodyStr), nil
}