第4个项目 - 实时聊天室 WebSocket

package models

import "time"

// MessageType 消息类型
type MessageType string

const (
	MessageTypeJoin    MessageType = "join"     // 加入聊天室
	MessageTypeLeave   MessageType = "leave"    // 离开聊天室
	MessageTypeChat    MessageType = "chat"     // 聊天消息
	MessageTypePrivate MessageType = "private"  // 私聊消息
	MessageTypeUsers   MessageType = "users"    // 用户列表
)

// Message 消息结构体
type Message struct {
	Type      MessageType `json:"type"`
	Username  string      `json:"username"`
	Content   string      `json:"content"`
	To        string      `json:"to,omitempty"`        // 私聊目标用户
	Timestamp time.Time   `json:"timestamp"`
	Users     []string    `json:"users,omitempty"`     // 在线用户列表
}

// NewMessage 创建新消息
func NewMessage(msgType MessageType, username, content string) *Message {
	return &Message{
		Type:      msgType,
		Username:  username,
		Content:   content,
		Timestamp: time.Now(),
	}
}

// NewPrivateMessage 创建私聊消息
func NewPrivateMessage(from, to, content string) *Message {
	return &Message{
		Type:      MessageTypePrivate,
		Username:  from,
		To:        to,
		Content:   content,
		Timestamp: time.Now(),
	}
}

package client

import (
	"encoding/json"
	"log"
	"time"

	"github.com/gorilla/websocket"
	"github.com/yourusername/chat-room/models"
)

const (
	// 写入超时时间
	writeWait = 10 * time.Second
	// 读取超时时间
	pongWait = 60 * time.Second
	// ping周期(必须小于pongWait)
	pingPeriod = (pongWait * 9) / 10
	// 最大消息大小
	maxMessageSize = 512
)

// Client 客户端连接
type Client struct {
	Hub      Hub
	Conn     *websocket.Conn
	Send     chan *models.Message
	Username string
}

// Hub接口(避免循环依赖)
type Hub interface {
	Register(client *Client)
	Unregister(client *Client)
	Broadcast(message *models.Message)
	SendPrivate(message *models.Message)
	GetOnlineUsers() []string
}

// ReadPump 从WebSocket读取消息
func (c *Client) ReadPump() {
	defer func() {
		c.Hub.Unregister(c)
		c.Conn.Close()
	}()

	c.Conn.SetReadDeadline(time.Now().Add(pongWait))
	c.Conn.SetPongHandler(func(string) error {
		c.Conn.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})

	for {
		_, messageData, err := c.Conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("错误: %v", err)
			}
			break
		}

		// 解析消息
		var msg models.Message
		if err := json.Unmarshal(messageData, &msg); err != nil {
			log.Printf("解析消息失败: %v", err)
			continue
		}

		// 设置用户名和时间戳
		msg.Username = c.Username
		msg.Timestamp = time.Now()

		// 根据消息类型处理
		switch msg.Type {
		case models.MessageTypeChat:
			// 广播消息
			c.Hub.Broadcast(&msg)
		case models.MessageTypePrivate:
			// 私聊消息
			c.Hub.SendPrivate(&msg)
		}
	}
}

// WritePump 向WebSocket写入消息
func (c *Client) WritePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.Conn.Close()
	}()

	for {
		select {
		case message, ok := <-c.Send:
			c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				// Hub关闭了通道
				c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			// 序列化消息
			messageData, err := json.Marshal(message)
			if err != nil {
				log.Printf("序列化消息失败: %v", err)
				continue
			}

			// 发送消息
			if err := c.Conn.WriteMessage(websocket.TextMessage, messageData); err != nil {
				return
			}

		case <-ticker.C:
			c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

package hub

import (
	"log"
	"sync"

	"github.com/yourusername/chat-room/client"
	"github.com/yourusername/chat-room/models"
)

// Hub 管理所有客户端连接
type Hub struct {
	// 已注册的客户端
	clients map[*client.Client]bool

	// 用户名到客户端的映射
	userClients map[string]*client.Client

	// 广播消息的通道
	broadcast chan *models.Message

	// 注册客户端的通道
	register chan *client.Client

	// 注销客户端的通道
	unregister chan *client.Client

	// 历史消息(最多保存100条)
	history []*models.Message

	// 互斥锁
	mu sync.RWMutex
}

// NewHub 创建新的Hub
func NewHub() *Hub {
	return &Hub{
		clients:     make(map[*client.Client]bool),
		userClients: make(map[string]*client.Client),
		broadcast:   make(chan *models.Message),
		register:    make(chan *client.Client),
		unregister:  make(chan *client.Client),
		history:     make([]*models.Message, 0, 100),
	}
}

// Run 启动Hub
func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			// 注册新客户端
			h.mu.Lock()
			h.clients[client] = true
			h.userClients[client.Username] = client
			h.mu.Unlock()

			// 发送历史消息给新用户
			h.sendHistory(client)

			// 广播用户加入消息
			joinMsg := models.NewMessage(models.MessageTypeJoin, client.Username, "加入了聊天室")
			h.addHistory(joinMsg)
			h.broadcastMessage(joinMsg)

			// 广播更新的用户列表
			h.broadcastUserList()

			log.Printf("用户 %s 加入聊天室, 当前在线: %d", client.Username, len(h.clients))

		case client := <-h.unregister:
			// 注销客户端
			if _, ok := h.clients[client]; ok {
				h.mu.Lock()
				delete(h.clients, client)
				delete(h.userClients, client.Username)
				close(client.Send)
				h.mu.Unlock()

				// 广播用户离开消息
				leaveMsg := models.NewMessage(models.MessageTypeLeave, client.Username, "离开了聊天室")
				h.addHistory(leaveMsg)
				h.broadcastMessage(leaveMsg)

				// 广播更新的用户列表
				h.broadcastUserList()

				log.Printf("用户 %s 离开聊天室, 当前在线: %d", client.Username, len(h.clients))
			}

		case message := <-h.broadcast:
			// 广播消息
			h.addHistory(message)
			h.broadcastMessage(message)
		}
	}
}

// Register 注册客户端
func (h *Hub) Register(c *client.Client) {
	h.register <- c
}

// Unregister 注销客户端
func (h *Hub) Unregister(c *client.Client) {
	h.unregister <- c
}

// Broadcast 广播消息
func (h *Hub) Broadcast(message *models.Message) {
	h.broadcast <- message
}

// SendPrivate 发送私聊消息
func (h *Hub) SendPrivate(message *models.Message) {
	h.mu.RLock()
	targetClient, exists := h.userClients[message.To]
	h.mu.RUnlock()

	if !exists {
		// 目标用户不在线
		return
	}

	// 发送给目标用户
	select {
	case targetClient.Send <- message:
	default:
		// 发送失败,关闭连接
		h.Unregister(targetClient)
	}
}

// GetOnlineUsers 获取在线用户列表
func (h *Hub) GetOnlineUsers() []string {
	h.mu.RLock()
	defer h.mu.RUnlock()

	users := make([]string, 0, len(h.clients))
	for client := range h.clients {
		users = append(users, client.Username)
	}
	return users
}

// broadcastMessage 广播消息给所有客户端
func (h *Hub) broadcastMessage(message *models.Message) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	for client := range h.clients {
		select {
		case client.Send <- message:
		default:
			// 发送失败,关闭连接
			close(client.Send)
			delete(h.clients, client)
			delete(h.userClients, client.Username)
		}
	}
}

// broadcastUserList 广播用户列表
func (h *Hub) broadcastUserList() {
	users := h.GetOnlineUsers()
	message := &models.Message{
		Type:  models.MessageTypeUsers,
		Users: users,
	}
	h.broadcastMessage(message)
}

// addHistory 添加到历史记录
func (h *Hub) addHistory(message *models.Message) {
	h.mu.Lock()
	defer h.mu.Unlock()

	h.history = append(h.history, message)
	// 只保留最近100条
	if len(h.history) > 100 {
		h.history = h.history[1:]
	}
}

// sendHistory 发送历史消息给新用户
func (h *Hub) sendHistory(c *client.Client) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	for _, msg := range h.history {
		select {
		case c.Send <- msg:
		default:
			return
		}
	}
}

package main

import (
	"log"
	"net/http"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/yourusername/chat-room/client"
	"github.com/yourusername/chat-room/hub"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	// 允许所有来源(生产环境应该限制)
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func main() {
	// 创建Hub并启动
	h := hub.NewHub()
	go h.Run()

	// 创建Gin引擎
	router := gin.Default()

	// 静态文件
	router.Static("/static", "./static")
	router.StaticFile("/", "./static/index.html")

	// WebSocket端点
	router.GET("/ws", func(c *gin.Context) {
		handleWebSocket(h, c.Writer, c.Request)
	})

	// 启动服务器
	log.Println("🚀 聊天室启动在 http://localhost:8080")
	if err := router.Run(":8080"); err != nil {
		log.Fatal(err)
	}
}

func handleWebSocket(h *hub.Hub, w http.ResponseWriter, r *http.Request) {
	// 获取用户名
	username := r.URL.Query().Get("username")
	if username == "" {
		http.Error(w, "缺少用户名", http.StatusBadRequest)
		return
	}

	// 升级到WebSocket
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("升级WebSocket失败: %v", err)
		return
	}

	// 创建客户端
	c := &client.Client{
		Hub:      h,
		Conn:     conn,
		Send:     make(chan *models.Message, 256),
		Username: username,
	}

	// 注册客户端
	h.Register(c)

	// 启动读写协程
	go c.WritePump()
	go c.ReadPump()
}

type ChatMessage struct {
	ID        uint      `gorm:"primaryKey"`
	Username  string    `gorm:"index"`
	Content   string
	CreatedAt time.Time
}

// 保存消息到数据库
db.Create(&ChatMessage{
	Username: msg.Username,
	Content:  msg.Content,
})

type Room struct {
	ID      string
	Hub     *Hub
	Clients map[*Client]bool
}

type Attachment struct {
	Type string `json:"type"` // image, file, emoji
	URL  string `json:"url"`
}