package main
import (
"fmt"
"github.com/gorilla/websocket"
"net/http"
)
var upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true // 允许跨域请求
},
}
// 模拟的消息推送中心
type PushCenter struct {
clients map[*websocket.Conn]bool
register chan *websocket.Conn
unregister chan *websocket.Conn
broadcast chan []byte
}
func NewPushCenter() *PushCenter {
return &PushCenter{
clients: make(map[*websocket.Conn]bool),
register: make(chan *websocket.Conn),
unregister: make(chan *websocket.Conn),
broadcast: make(chan []byte),
}
}
func (p *PushCenter) Run() {
for {
select {
case conn := <-p.register:
p.clients[conn] = true
case conn := <-p.unregister:
if _, ok := p.clients[conn]; ok {
delete(p.clients, conn)
conn.Close()
}
case message := <-p.broadcast:
for conn := range p.clients {
err := conn.WriteMessage(websocket.TextMessage, message)
if err != nil {
p.unregister <- conn
}
}
}
}
}
func (p *PushCenter) ServeWs(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
fmt.Println(err)
return
}
p.register <- conn
}
func main() {
pushCenter := NewPushCenter()
go pushCenter.Run()
http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
pushCenter.ServeWs(w, r)
})
http.ListenAndServe(":8080", nil)
}这段代码实现了一个简单的WebSocket消息推送中心。它创建了一个PushCenter结构体，用于管理所有的WebSocket连接。ServeWs函数处理WebSocket连接请求，并将新的连接加入到注册通道中。Run方法是消息推送中心的核心，它接收和处理三种类型的通道消息：注册新连接、注销旧连接以及广播消息。这个简单的实现展示了如何使用Go语言和gorilla/websocket库来创建一个基本的WebSocket服务器。