feat: update

api/chatroom/v1/chatroom.proto

@@ -7,9 +7,8 @@ option go_package = "api/chatroom/v1;v1";
 service ChatRoom {
 }
 
-message WebsocketProto  {
-  string event_id = 1;
-  string payload = 2;
+enum MessageType {
+  Chat = 0;
 }
 
 message ChatMessage {

app/chatroom/internal/server/registrar.go

@@ -4,6 +4,7 @@ import (
 	"github.com/go-kratos/kratos/contrib/registry/consul/v2"
 	"github.com/go-kratos/kratos/v2/registry"
 	consulAPI "github.com/hashicorp/consul/api"
+
 	"kratos-chatroom/app/chatroom/internal/conf"
 )
 

app/chatroom/internal/server/websocket.go

@@ -1,22 +1,39 @@
 package server
 
 import (
+	"errors"
+
+	"github.com/go-kratos/kratos/v2/encoding"
 	"github.com/go-kratos/kratos/v2/log"
 	"github.com/tx7do/kratos-transport/transport/websocket"
 
+	v1 "kratos-chatroom/api/chatroom/v1"
 	"kratos-chatroom/app/chatroom/internal/conf"
 	"kratos-chatroom/app/chatroom/internal/service"
 )
 
 // NewWebsocketServer create a websocket server.
 func NewWebsocketServer(c *conf.Server, _ log.Logger, svc *service.ChatRoomService) *websocket.Server {
 	srv := websocket.NewServer(
-		websocket.Address(c.Websocket.Addr),
-		websocket.ReadHandle(c.Websocket.Path, svc.OnWebsocketMessage),
-		websocket.ConnectHandle(svc.OnWebsocketConnect),
+		websocket.WithAddress(c.Websocket.Addr),
+		websocket.WithPath(c.Websocket.Path),
+		websocket.WithConnectHandle(svc.OnWebsocketConnect),
+		websocket.WithCodec(encoding.GetCodec("json")),
 	)
 
 	svc.SetWebsocketServer(srv)
 
+	srv.RegisterMessageHandler(websocket.MessageType(v1.MessageType_Chat),
+		func(sessionId websocket.SessionID, payload websocket.MessagePayload) error {
+			switch t := payload.(type) {
+			case *v1.ChatMessage:
+				return svc.OnChatMessage(sessionId, t)
+			default:
+				return errors.New("invalid payload type")
+			}
+		},
+		func() websocket.Any { return &v1.ChatMessage{} },
+	)
+
 	return srv
 }

app/chatroom/internal/service/websocket.go

@@ -1,8 +1,6 @@
 package service
 
 import (
-	"encoding/json"
-	"fmt"
 	"github.com/tx7do/kratos-transport/transport/websocket"
 	v1 "kratos-chatroom/api/chatroom/v1"
 )
@@ -11,54 +9,15 @@ func (s *ChatRoomService) SetWebsocketServer(ws *websocket.Server) {
 	s.ws = ws
 }
 
-func (s *ChatRoomService) OnWebsocketMessage(connectionId string, message *websocket.Message) error {
-	s.log.Infof("[%s] Payload: %s\n", connectionId, string(message.Body))
-
-	var proto v1.WebsocketProto
-
-	if err := json.Unmarshal(message.Body, &proto); err != nil {
-		s.log.Error("Error unmarshalling proto json %v", err)
-		return nil
-	}
-
-	switch proto.EventId {
-	case "chat":
-		chatMsg := proto.Payload
-		fmt.Println("chat message:", chatMsg)
-		_ = s.OnChatMessage(connectionId, &chatMsg)
-	}
-
-	return nil
-}
-
-func (s *ChatRoomService) OnChatMessage(connectionId string, msg *string) error {
-	s.BroadcastToWebsocketClient("chat", msg)
-	return nil
-}
-
-func (s *ChatRoomService) OnWebsocketConnect(connectionId string, register bool) {
+func (s *ChatRoomService) OnWebsocketConnect(sessionId websocket.SessionID, register bool) {
 	if register {
-		fmt.Printf("%s connected\n", connectionId)
+		s.log.Infof("%s connected\n", sessionId)
 	} else {
-		fmt.Printf("%s disconnect\n", connectionId)
+		s.log.Infof("%s disconnect\n", sessionId)
 	}
 }
 
-func (s *ChatRoomService) BroadcastToWebsocketClient(eventId string, payload interface{}) {
-	if payload == nil {
-		return
-	}
-
-	bufPayload, _ := json.Marshal(&payload)
-
-	var proto v1.WebsocketProto
-	proto.EventId = eventId
-	proto.Payload = string(bufPayload)
-
-	bufProto, _ := json.Marshal(&proto)
-
-	var msg websocket.Message
-	msg.Body = bufProto
-
-	s.ws.Broadcast(&msg)
+func (s *ChatRoomService) OnChatMessage(sessionId websocket.SessionID, msg *v1.ChatMessage) error {
+	s.ws.Broadcast(websocket.MessageType(v1.MessageType_Chat), msg)
+	return nil
 }