- 安装必要的依赖包
go get github.com/gorilla/websocket
- 创建一个websocket服务器
package main
import (
"fmt"
"net/http"
"github.com/gorilla/websocket"
)
// 定义一个全局的websocket升级器
var upgrader = websocket.Upgrader{
// 允许跨域
CheckOrigin: func(r *http.Request) bool {
return true
},
}
// 定义一个全局的map,用来存储在线用户
var onlineUsers = make(map[*websocket.Conn]string)
// 定义一个全局的channel,用来接收用户发送的消息
var message = make(chan []byte)
func main() {
// 创建一个http服务
http.HandleFunc("/ws", wsHandler)
http.ListenAndServe("0.0.0.0:7777", nil)
}
// 定义一个处理websocket请求的函数
func wsHandler(w http.ResponseWriter, r *http.Request) {
// 将http协议升级为websocket协议
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
fmt.Println("upgrade failed, err:", err)
return
}
// 将新接入的用户添加到onlineUsers中
onlineUsers[conn] = ""
// 处理用户发送的消息
go processMessage(conn)
// 处理用户断开连接
go processDisconnect(conn)
}
- 处理用户发送的消息
// 处理用户发送的消息
func processMessage(conn *websocket.Conn) {
for {
// 读取用户发送的消息
_, data, err := conn.ReadMessage()
if err != nil {
fmt.Println("read message failed, err:", err)
return
}
// 将消息发送到message channel中
message <- data
}
}
- 处理用户断开连接
// 处理用户断开连接
func processDisconnect(conn *websocket.Conn) {
for {
// 读取用户发送的消息
_, _, err := conn.ReadMessage()
if err != nil {
// 将断开的用户从onlineUsers中删除
delete(onlineUsers, conn)
fmt.Println("disconnect success")
return
}
}
}
- 广播消息
// 广播消息
func broadcastMessage() {
for {
// 从message channel中读取消息
data := <-message
// 遍历onlineUsers,将消息发送给每一个在线用户
for conn, _ := range onlineUsers {
conn.WriteMessage(websocket.TextMessage, data)
}
}
}
- 运行
func main() {
// 创建一个http服务
http.HandleFunc("/ws", wsHandler)
http.ListenAndServe("0.0.0.0:7777", nil)
// 启动一个goroutine,用来广播消息
go broadcastMessage()
}