transfer-site/internal/http/handlers/ws.go

package handlers

import (
	"encoding/json"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		// 允许所有来源(生产环境应该做更严格的检查)
		return true
	},
}

// ProgressMessage WebSocket进度消息结构
type ProgressMessage struct {
	Status  string `json:"status"`  // 状态: checking, cloning, pulling, success, error
	Message string `json:"message"` // 消息文本
	Percent int    `json:"percent"` // 进度百分比 (0-100)
}

// ProgressHub 进度广播中心
type ProgressHub struct {
	clients    map[string]*websocket.Conn
	clientsMux sync.RWMutex
	broadcast  chan ProgressMessage
}

var (
	hub     *ProgressHub
	hubOnce sync.Once
)

// GetProgressHub 获取进度广播中心单例
func GetProgressHub() *ProgressHub {
	hubOnce.Do(func() {
		hub = &ProgressHub{
			clients:   make(map[string]*websocket.Conn),
			broadcast: make(chan ProgressMessage, 10),
		}
		go hub.run()
	})
	return hub
}

// run 运行广播循环
func (h *ProgressHub) run() {
	for msg := range h.broadcast {
		h.clientsMux.RLock()
		for clientID, conn := range h.clients {
			err := conn.WriteJSON(msg)
			if err != nil {
				// 写入失败,移除客户端
				conn.Close()
				h.clientsMux.RUnlock()
				h.removeClient(clientID)
				h.clientsMux.RLock()
			}
		}
		h.clientsMux.RUnlock()
	}
}

// addClient 添加客户端
func (h *ProgressHub) addClient(clientID string, conn *websocket.Conn) {
	h.clientsMux.Lock()
	defer h.clientsMux.Unlock()
	h.clients[clientID] = conn
}

// removeClient 移除客户端
func (h *ProgressHub) removeClient(clientID string) {
	h.clientsMux.Lock()
	defer h.clientsMux.Unlock()
	delete(h.clients, clientID)
}

// SendProgress 发送进度消息
func (h *ProgressHub) SendProgress(msg ProgressMessage) {
	select {
	case h.broadcast <- msg:
	case <-time.After(1 * time.Second):
		// 超时,丢弃消息
	}
}

// GitProgressWS WebSocket处理器
func GitProgressWS() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// 升级HTTP连接到WebSocket
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			http.Error(w, "Failed to upgrade to WebSocket", http.StatusInternalServerError)
			return
		}
		defer conn.Close()

		// 生成客户端ID
		clientID := r.RemoteAddr + "-" + time.Now().Format("20060102150405")

		// 添加到hub
		hub := GetProgressHub()
		hub.addClient(clientID, conn)
		defer hub.removeClient(clientID)

		// 发送欢迎消息
		welcomeMsg := ProgressMessage{
			Status:  "connected",
			Message: "WebSocket连接已建立",
			Percent: 0,
		}
		if err := conn.WriteJSON(welcomeMsg); err != nil {
			return
		}

		// 保持连接,读取客户端消息(主要用于心跳)
		for {
			_, _, err := conn.ReadMessage()
			if err != nil {
				// 连接断开
				break
			}
		}
	}
}

// SendProgressToClients 辅助函数:向所有客户端发送进度
func SendProgressToClients(status, message string, percent int) {
	hub := GetProgressHub()
	hub.SendProgress(ProgressMessage{
		Status:  status,
		Message: message,
		Percent: percent,
	})
}

// SendProgressJSON 辅助函数:发送JSON格式的进度消息
func SendProgressJSON(w http.ResponseWriter, status, message string, percent int) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(ProgressMessage{
		Status:  status,
		Message: message,
		Percent: percent,
	})
}