websocket路由封装示例

ws或者tcp不像http有路由，只能自定义消息格式，拆出来路由

message

package messageimport ("encoding/json""github.com/olahol/melody"
)// Message 消息结构体
type Message struct {Router string      `json:"router"`Msg    interface{} `json:"msg"`
}// UserLoginMessage 用户登录消息结构体
type UserLoginMessage struct {Username string `json:"username"`Password string `json:"password"`
}// UserRegisterMessage 用户注册消息结构体
type UserRegisterMessage struct {Username string `json:"username"`Password string `json:"password"`Email    string `json:"email"`
}// RouteNotFoundMessage 路由找不到消息结构体
type RouteNotFoundMessage struct {Router string `json:"router"`
}// ErrorMessage 错误消息结构体
type ErrorMessage struct {Code    string `json:"code"`Message string `json:"message"`
}// ParseMessageToStruct 将Message结构体中的Msg字段解析到目标结构体
func ParseMessageToStruct(msg *Message, target interface{}) error {// 将interface{}转换为JSON字节msgBytes, err := json.Marshal(msg.Msg)if err != nil {return err}// 将JSON字节解析到目标结构体return json.Unmarshal(msgBytes, target)
}// NewErrorMessage 创建新的错误消息
func NewErrorMessage(code, message string) *Message {return &Message{Router: "error",Msg: ErrorMessage{Code:    code,Message: message,},}
}// SendErrorMessage 发送错误消息到客户端
func SendErrorMessage(s *melody.Session, msg *Message) error {// 将消息转换为JSON字节msgBytes, err := json.Marshal(msg)if err != nil {return err}// 发送消息return s.Write(msgBytes)
}// NewRouteNotFoundMessage 创建路由找不到消息
func NewRouteNotFoundMessage(router string) *Message {return &Message{Router: "route.not.found",Msg: RouteNotFoundMessage{Router: router,},}
}

router

package routerimport ("app01/message""log""sync""github.com/olahol/melody"
)// HandlerFunc 定义消息处理函数类型
type HandlerFunc func(*message.Message, *melody.Session)// IRouter 定义路由处理器接口
type IRouter interface {// Handle 处理消息Handle(*message.Message, *melody.Session)
}// 适配器结构，用于将HandlerFunc转换为IRouter
type handlerFuncAdapter struct {handler HandlerFunc
}// Handle 实现IRouter接口
func (a *handlerFuncAdapter) Handle(msg *message.Message, s *melody.Session) {a.handler(msg, s)
}// RouterManager 路由管理器
type RouterManager struct {routers map[string]IRoutermutex   sync.RWMutex
}var (// GlobalRouterManager 全局路由管理器实例GlobalRouterManager *RouterManager
)// 初始化路由管理器
func init() {GlobalRouterManager = &RouterManager{routers: make(map[string]IRouter),}
}// RegisterRouter 注册路由处理器(接口版本)
func RegisterRouter(routerName string, router IRouter) {GlobalRouterManager.mutex.Lock()defer GlobalRouterManager.mutex.Unlock()GlobalRouterManager.routers[routerName] = router
}// RegisterHandlerFunc 注册路由处理器(函数版本)
func RegisterHandlerFunc(routerName string, handler HandlerFunc) {RegisterRouter(routerName, &handlerFuncAdapter{handler: handler})
}// HandleMessage 处理消息
func HandleMessage(msg *message.Message, s *melody.Session) {GlobalRouterManager.mutex.RLock()router, ok := GlobalRouterManager.routers[msg.Router]GlobalRouterManager.mutex.RUnlock()if ok {router.Handle(msg, s)} else {// 路由找不到的处理逻辑HandleNotFoundRoute(msg, s)}
}// HandleNotFoundRoute 处理路由找不到的情况
func HandleNotFoundRoute(msg *message.Message, s *melody.Session) {log.Printf("路由找不到: %s\n", msg.Router)// 创建路由找不到消息notFoundMsg := message.NewRouteNotFoundMessage(msg.Router)// 发送消息给客户端message.SendErrorMessage(s, notFoundMsg)
}

user_router 后面可以根据业务扩展多个

package routerimport ("app01/message""log""github.com/olahol/melody"
)// 方式1: 使用闭包注册用户路由
func init() {// 使用匿名函数注册用户登录消息处理器RegisterHandlerFunc("user.login", func(msg *message.Message, s *melody.Session) {// 解析JSON到结构体var loginMsg message.UserLoginMessageif err := message.ParseMessageToStruct(msg, &loginMsg); err != nil {log.Printf("解析登录消息失败: %v\n", err)// 发送错误消息errMsg := message.NewErrorMessage("登录参数错误", err.Error())message.SendErrorMessage(s, errMsg)return}// 简单验证if loginMsg.Username != "" && loginMsg.Password != "" {log.Printf("用户 %s 登录成功\n", loginMsg.Username)// 这里可以添加实际的登录逻辑} else {log.Printf("用户名或密码不能为空\n")// 发送错误消息errMsg := message.NewErrorMessage("登录失败", "用户名或密码不能为空")message.SendErrorMessage(s, errMsg)}})// 注册用户注册消息处理器RegisterHandlerFunc("user.register", handleRegisterMessage)
}// 命名处理函数
func handleRegisterMessage(msg *message.Message, s *melody.Session) {// 解析JSON到结构体var registerMsg message.UserRegisterMessageif err := message.ParseMessageToStruct(msg, &registerMsg); err != nil {log.Printf("解析注册消息失败: %v\n", err)// 发送错误消息errMsg := message.NewErrorMessage("注册参数错误", err.Error())message.SendErrorMessage(s, errMsg)return}// 简单验证if registerMsg.Username != "" && registerMsg.Password != "" && registerMsg.Email != "" {log.Printf("用户 %s 注册成功\n", registerMsg.Username)// 这里可以添加实际的注册逻辑} else {log.Printf("用户名、密码或邮箱不能为空\n")// 发送错误消息errMsg := message.NewErrorMessage("注册失败", "用户名、密码或邮箱不能为空")message.SendErrorMessage(s, errMsg)}
}