WebSocket原来还能这么玩

原创

闫同学

发布于 2024-04-21 09:06:45

2961

WebSocket是一种在单个TCP连接上进行全双工通信的协议。该协议于2011年被IETF定为标准RFC 6455，并由RFC7936补充规范。WebSocket API也被W3C定为标准。这种通信方式使得客户端和服务器之间的数据交换变得更加简单，允许服务端主动向客户端推送数据。在WebSocket API中，客户端和服务器只需要完成一次握手，之后两者之间就可以直接创建持久性的连接，并进行双向数据传输。

这篇文章我就准备使用不到200行Go代码使用WebSocket实现一个简单的私聊系统，正文开始~

1 环境准备

首先，我们需要安装gorilla/websocket包，它是Go语言中一个非常流行的WebSocket库：

go get -u github.com/gorilla/websocket

2 功能设计

其实发送私聊消息无非就是解决两个问题：

给谁发
发什么

延伸一下，在我们的功能设计中就分为两步，一是确认身份，就需要客户端有一个向服务端注册的动作，二是发消息，而且在发消息的同时除了内容以外，还要知道发送者和接受者都是谁。

2.1 用户注册

作为私聊系统的第一步，客户端需要先与服务端建立WebSocket连接，抽象出User和用户连接池，在注册时发送消息然后由客户端保存该User对应的WebSocket连接。

2.2 消息发送

消息发送时要使用固定的结构体，使用JSON格式进行序列化，首先是客户端进行消息的构建，然后发送到服务器，此时服务器就类似于一个路由器，将之前保存的接收方的WebSocket连接拿出来进行消息的发送。

3 代码实现

下面我们进行代码的实现。

3.1 消息定义

首先是结构体，也就是消息类型的定义，Chat作为客户端和服务器交互的统一类型，承载了事件类型（注册或发送消息事件）和消息体，消息体中有接受者、发送者和消息内容等主要字段。

package model

type Event int

const (
	Register Event = 1
	SendMsg        = 2
)

type Chat struct {
	Event   Event
	Message Message
	User    User
}

type Message struct {
	SendUser   *User
	Receiver   *User
	Content    string
	CreateTime string
}

type User struct {
	UserId   int64
	UserName string
	Address  string
}

3.2 服务端代码

服务端代码主要是提供一个WebSocket连接，根据接收的消息事件取出相关的消息，再进行后续的注册或发送消息逻辑。

var (
	ws      = websocket.Upgrader{}
	userMap map[int64]*websocket.Conn
)

func init() {
	userMap = make(map[int64]*websocket.Conn)
}

func main() {
	http.HandleFunc("/privateChat", privateChat)
	_ = http.ListenAndServe(":9900", nil)
}

func privateChat(w http.ResponseWriter, r *http.Request) {
	c, err := ws.Upgrade(w, r, nil) //升级将 HTTP 服务器连接升级到 WebSocket 协议
	if err != nil {
		log.Printf("upgrade err:%s\n", err)
		return
	}

	for {
		mt, message, err := c.ReadMessage()
		if err != nil {
			log.Printf("read message err:%s\n", err)
			continue
		}

		var chat model.Chat
		if err := json.Unmarshal(message, &chat); err != nil {
			log.Printf("unmarshal err:%s\n", err)
			continue
		}

		switch chat.Event {
		case model.Register:
			//设置userID
			chat.User.UserId = time.Now().Unix()
			//保存用户连接
			userMap[chat.User.UserId] = c
			break
		case model.SendMsg:
			chat.Message.CreateTime = time.Now().Format("2006-01-02 15:04:05")
			//拿到用户连接
			ok := false
			c, ok = userMap[chat.Message.Receiver.UserId]
			if !ok {
				//如果没有，拿到发送方用户的连接，告诉他不行
				c, _ = userMap[chat.Message.SendUser.UserId]
				chat.Message.Receiver = chat.Message.SendUser
				chat.Message.Content = "发送失败"
			}
			break
		default:
			c, _ = userMap[chat.Message.SendUser.UserId]
			chat.Message.Receiver = chat.Message.SendUser
			chat.Message.Content = "消息类型不对"
		}

		log.Printf("now chat : %+v \n", chat)

		//响应数据
		bytes, err := json.Marshal(chat)
		if err != nil {
			log.Printf("marshal err:%s\n", err)
			continue
		}
		if err := c.WriteMessage(mt, bytes); err != nil {
			log.Printf("write message err:%s\n", err)
			continue
		}
	}
}

3.3 客户端代码

客户端代码首先是连接WebSocket服务端，然后再启动一个HTTP服务用于接收和发送具体的WebSocket消息。

const (
	PrivateChatUrl = "ws://127.0.0.1:9900/privateChat"
)

var (
	user model.User
	name string
	port int 
)

func init() {
	flag.StringVar(&name, "name", "", "user name")
	flag.IntVar(&port, "port", 8801, "server port")
}

func main() {
	flag.Parse()

	c, _, err := websocket.DefaultDialer.Dial(PrivateChatUrl, nil)
	if err != nil {
		log.Fatal("dial:", err)
		return
	}

	//注册
	u := model.Chat{Event: model.Register, User: model.User{UserName: name}}
	bytes, _ := json.Marshal(u)
	_ = c.WriteMessage(websocket.TextMessage, bytes)

	//读取消息监听
	go func() {
		for {
			_, msg, err := c.ReadMessage()
			if err != nil {
				log.Printf("read message err:%s\n", err)
				continue
			}
			var res model.Chat
			if err := json.Unmarshal(msg, &res); err != nil {
				log.Printf("unmarshal err:%s\n", err)
				continue
			}
			switch res.Event {
			case model.Register:
				user = res.User
				fmt.Printf("register success , now user %v \n", user)
			case model.SendMsg:
				fmt.Printf("用户 %s 在 %s 给你发送了一条消息:%s \n",
					res.Message.SendUser.UserName, res.Message.CreateTime, res.Message.Content)
			}
		}
	}()

	http.HandleFunc("/send", func(w http.ResponseWriter, r *http.Request) {
		uid := r.URL.Query().Get("uid")
		content := r.URL.Query().Get("content")
		msg, _ := json.Marshal(model.Chat{
			Event: model.SendMsg,
			Message: model.Message{
				SendUser: &user,
				Receiver: &model.User{UserId: cast.ToInt64(uid)},
				Content:  content,
			},
		})
		_ = c.WriteMessage(websocket.TextMessage, msg)
		_, _ = w.Write([]byte("ok"))
	})
	_ = http.ListenAndServe(fmt.Sprintf(":%d",port), nil)
}

4 使用方式

首先启动服务端：

go run server/main.go

然后启动客户端，命名为ZhangSan，HTTP端口号为8801，进行注册：

go run client/main.go -port 8801 -name ZhangSan

得到的相应结果：

register success , now user {1713606616 ZhangSan }

然后启动另一个客户端，命名为LiSi，HTTP端口号为8802，进行注册：

go run client/main.go -port 8802 -name LiSi

得到的相应结果：

register success , now user {1713606595 LiSi }

现在我们知道了两个用户的uId，然后我们可以进行发送消息：

上图是用户LiSi给用户ZhangSan发送消息，我们查看接收方的响应：

完成~

5 小总结

不知道上面的代码大家有没有看出问题，没错，Go语言的map类型是线程不安全的，因此最好进行在操作时进行加锁。

//Lock  
userMap[userId] = WebSocketConn //具体操作
//Unlock

除此之外，还有许多功能可以扩展，希望大家能给出建议~

我正在参与2024腾讯技术创作特训营最新征文，快来和我瓜分大奖！

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

腾讯技术创作特训营S6

websocket

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

登录后参与评论

0 条评论

热度