Go使用WebSocket实现一个公域聊天室

var (
	// 存储所有在线的WebSocket连接（核心：全员广播的基础）
	clients = make(map[*websocket.Conn]bool)
	// 存储连接对应的用户名（区分谁发的消息）
	clientNames = make(map[*websocket.Conn]string)
	// 互斥锁：多个协程同时修改map，必须加锁，否则程序崩溃
	clientsMutex sync.Mutex
	// HTTP 升级为 WebSocket 的工具
	upgrader = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true // 允许跨域（浏览器访问必备）
		},
	}
)

func main() {
	// 注册路由：访问 /ws 就触发聊天处理逻辑
	http.HandleFunc("/ws", wsHandler)

	fmt.Println("WebSocket 服务启动：ws://127.0.0.1:8081/ws")
	// 核心：启动HTTP服务，自动为每个客户端开goroutine（不阻塞）
	if err := http.ListenAndServe(":8081", nil); err != nil {
		fmt.Println("服务启动失败：", err)
	}
}

// 把普通HTTP请求，升级为长连接WebSocket
conn, err := upgrader.Upgrade(w, r, nil)
// 函数结束自动关闭连接（用户断开时触发）
defer conn.Close()

// 服务端主动发消息：欢迎语
conn.WriteMessage(websocket.TextMessage, []byte("欢迎连接！请输入你的名字："))
// 读取用户输入的名字
_, msg, err := conn.ReadMessage()
name := string(msg)

// 把当前用户加入在线列表
registerClient(conn, name)
// 广播：XXX 进入聊天室（全员可见）
broadcast("系统: " + name + " 进入了聊天室")

// defer：函数最后执行（用户关闭网页时触发）
defer func() {
	unregisterClient(conn) // 从在线列表删除
	broadcast("系统: " + name + " 离开了聊天室") // 广播离开消息
}()

// 死循环：持续监听用户发的消息
for {
	_, msg, err := conn.ReadMessage()
	if err != nil {
		break // 出错/断开，退出循环
	}
	// 拼接消息：用户名: 内容
	chatMsg := fmt.Sprintf("%s: %s", name, string(msg))
	// 广播给所有人
	broadcast(chatMsg)
}

// 注册用户：加锁 → 修改 → 解锁
func registerClient(conn *websocket.Conn, name string) {
	clientsMutex.Lock()   // 加锁
	defer clientsMutex.Unlock() // 自动解锁
	clients[conn] = true
	clientNames[conn] = name
}

// 注销用户：从map删除
func unregisterClient(conn *websocket.Conn) {
	clientsMutex.Lock()
	defer clientsMutex.Unlock()
	delete(clients, conn)
	delete(clientNames, conn)
}

func broadcast(message string) {
	clientsMutex.Lock()
	defer clientsMutex.Unlock()

	// 遍历所有在线用户
	for client := range clients {
		// 给每个用户发送消息
		client.WriteMessage(websocket.TextMessage, []byte(message))
	}
}

package main

import (
    "fmt"
    "net/http"
    "sync"

    "github.com/gorilla/websocket" // go get github.com/gorilla/websocket
)

var (
    // clients 存储所有在线的 WebSocket 连接
    clients = make(map[*websocket.Conn]bool)
    // clientNames 存储连接对应的用户名
    clientNames = make(map[*websocket.Conn]string)
    // clientsMutex 保护 clients 和 clientNames 的并发读写
    clientsMutex sync.Mutex
    // upgrader 用于将 HTTP 请求升级为 WebSocket 请求
    upgrader = websocket.Upgrader{
       CheckOrigin: func(r *http.Request) bool {
          return true // 允许跨域
       },
    }
)

func main() {
    // 注册 WebSocket 路由
    http.HandleFunc("/ws", wsHandler)

    fmt.Println("WebSocket 服务启动：ws://127.0.0.1:8081/ws")
    if err := http.ListenAndServe(":8081", nil); err != nil {
       fmt.Println("服务启动失败：", err)
    }
}

// wsHandler 处理 WebSocket 连接请求
func wsHandler(w http.ResponseWriter, r *http.Request) {
    // 1. 升级 HTTP 连接为 WebSocket 连接
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
       fmt.Println("升级失败：", err)
       return
    }
    // 确保连接最终被关闭
    defer conn.Close()

    // 2. 询问并读取用户名
    if err := conn.WriteMessage(websocket.TextMessage, []byte("欢迎连接！请输入你的名字：")); err != nil {
       return
    }
    _, msg, err := conn.ReadMessage()
    if err != nil {
       return
    }
    name := string(msg)

    // 3. 注册新用户并广播进入消息
    registerClient(conn, name)
    broadcast("系统: " + name + " 进入了聊天室")
    fmt.Printf("用户 [%s] 已加入\n", name)

    // 4. 设置断开连接时的清理操作
    defer func() {
       unregisterClient(conn)
       broadcast("系统: " + name + " 离开了聊天室")
       fmt.Printf("用户 [%s] 已断开\n", name)
    }()

    // 5. 循环读取消息并广播
    for {
       _, msg, err := conn.ReadMessage()
       if err != nil {
          break // 读取错误或连接关闭
       }

       // 构造并广播消息
       chatMsg := fmt.Sprintf("%s: %s", name, string(msg))
       fmt.Println(chatMsg) // 服务端日志
       broadcast(chatMsg)
    }
}

// registerClient 线程安全地注册新连接
func registerClient(conn *websocket.Conn, name string) {
    clientsMutex.Lock()
    defer clientsMutex.Unlock()

    clients[conn] = true
    clientNames[conn] = name
}

// unregisterClient 线程安全地注销连接
func unregisterClient(conn *websocket.Conn) {
    clientsMutex.Lock()
    defer clientsMutex.Unlock()

    delete(clients, conn)
    delete(clientNames, conn)
}

// broadcast 向所有在线客户端广播消息
func broadcast(message string) {
    clientsMutex.Lock()
    defer clientsMutex.Unlock()

    msgBytes := []byte(message)
    for client := range clients {
       // 注意：如果某个客户端网络阻塞，这里可能会阻塞较长时间
       // 实际生产中应将发送逻辑放入独立的 goroutine 或使用 channel
       err := client.WriteMessage(websocket.TextMessage, msgBytes)
       if err != nil {
          fmt.Printf("广播消息失败: %v\n", err)
          client.Close()
          delete(clients, client)
          delete(clientNames, client)
       }
    }
}