web/server/pkg/weblive/danmaku.go

package weblive

import (
	"encoding/json"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"yh_web/server/handlers"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

const maxDanmakuRunes = 120

var allowedGifts = map[string]struct{}{
	"rocket":     {},
	"sports_car": {},
	"plane":      {},
	"carnival":   {},
	"rose":       {},
	"heart":      {},
	"star":       {},
	"clap":       {},
	"cake":       {},
	"crown":      {},
	"fireworks":  {},
	"gift_box":   {},
	"beer":       {},
	"mic":        {},
}

var (
	danmakuClientsMu sync.Mutex
	danmakuClients   = make(map[*websocket.Conn]string)
)

func writeDanmakuJSON(ws *websocket.Conn, v any) error {
	b, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return ws.WriteMessage(websocket.TextMessage, b)
}

func clipDanmakuText(t string) string {
	t = strings.TrimSpace(t)
	if t == "" {
		return ""
	}
	if utf8.RuneCountInString(t) <= maxDanmakuRunes {
		return t
	}
	runes := []rune(t)
	return string(runes[:maxDanmakuRunes])
}

// handleDanmakuWS 弹幕 / 礼物：未带有效 token 仅可收广播。礼物 JSON {"type":"gift","gift":"rocket"}；弹幕 {"text":"..."}
func handleDanmakuWS(c *gin.Context) {
	claims, tokenOK := handlers.ParseSiteClaims(c.Query("token"))
	canSend := tokenOK
	fromDisplay := "***"
	if tokenOK && claims != nil {
		fromDisplay = handlers.MaskSiteUsernameForDanmaku(claims.Username)
	}

	ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}
	ws.SetReadLimit(4096)
	clientIP := c.ClientIP()
	sessionID := RegisterOnlineSession("danmaku", clientIP, fromDisplay)

	danmakuClientsMu.Lock()
	danmakuClients[ws] = clientIP
	danmakuClientsMu.Unlock()

	defer func() {
		UnregisterOnlineSession(sessionID)
		danmakuClientsMu.Lock()
		delete(danmakuClients, ws)
		danmakuClientsMu.Unlock()
		_ = ws.Close()
	}()

	for {
		mt, payload, err := ws.ReadMessage()
		if err != nil {
			return
		}
		if mt != websocket.TextMessage {
			continue
		}
		TouchOnlineSession(sessionID)
		if IsMutedForIP(clientIP) {
			_ = writeDanmakuJSON(ws, map[string]interface{}{
				"type":    "error",
				"code":    "muted",
				"message": "当前已被禁言",
			})
			continue
		}
		if !AllowSendByIP(clientIP) {
			_ = writeDanmakuJSON(ws, map[string]interface{}{
				"type":    "error",
				"code":    "rate_limited",
				"message": "同 IP 发送过快，请稍后再试",
			})
			continue
		}
		if !canSend {
			_ = writeDanmakuJSON(ws, map[string]interface{}{
				"type":    "error",
				"code":    "login_required",
				"message": "请先登录或注册后再发弹幕或礼物",
			})
			continue
		}
		var envelope struct {
			Type string `json:"type"`
			Text string `json:"text"`
			Gift string `json:"gift"`
		}
		if err := json.Unmarshal(payload, &envelope); err != nil {
			continue
		}
		if strings.EqualFold(strings.TrimSpace(envelope.Type), "gift") {
			gid := strings.TrimSpace(envelope.Gift)
			if _, ok := allowedGifts[gid]; !ok {
				_ = writeDanmakuJSON(ws, map[string]interface{}{
					"type":    "error",
					"code":    "bad_gift",
					"message": "无效的礼物",
				})
				continue
			}
			out, err := json.Marshal(map[string]interface{}{
				"type": "gift",
				"gift": gid,
				"from": fromDisplay,
				"ts":   time.Now().UnixMilli(),
			})
			if err != nil {
				continue
			}
			danmakuBroadcast(out)
			continue
		}
		text := clipDanmakuText(envelope.Text)
		if text == "" {
			continue
		}
		out, err := json.Marshal(map[string]interface{}{
			"type": "dm",
			"text": text,
			"from": fromDisplay,
			"ts":   time.Now().UnixMilli(),
		})
		if err != nil {
			continue
		}
		danmakuBroadcast(out)
	}
}

func danmakuBroadcast(b []byte) {
	danmakuClientsMu.Lock()
	defer danmakuClientsMu.Unlock()
	dead := make([]*websocket.Conn, 0)
	for conn := range danmakuClients {
		_ = conn.SetWriteDeadline(time.Now().Add(8 * time.Second))
		if err := conn.WriteMessage(websocket.TextMessage, b); err != nil {
			dead = append(dead, conn)
		}
	}
	for _, conn := range dead {
		delete(danmakuClients, conn)
		_ = conn.Close()
	}
}