soul-yongping/soul-api/internal/handler/ws.go

// Package handler - WebSocket 占位：用户在线检测
// 小程序连接 WSS 发心跳，Redis 记录在线；管理端通过 HTTP 获取在线人数
// 后续可扩展：管理端 WSS 订阅、消息推送等

package handler

import (
	"context"
	"encoding/json"
	"log"
	"net/http"
	"strings"
	"time"

	"soul-api/internal/config"
	"soul-api/internal/database"
	"soul-api/internal/model"
	"soul-api/internal/redis"

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

const (
	wsOnlinePrefix   = "user:online:"
	wsOfflineTimeout = 300 // 5 分钟无心跳视为离线（秒）
)

var wsUpgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// WsMiniprogram 处理小程序 WSS 连接：鉴权后记录心跳，维持在线状态
// 路径：GET /ws/miniprogram?token=xxx
// 首条消息需包含 {"type":"auth","userId":"user_xxx"}，占位阶段不校验 token
// 容错：panic 时 recover 并关闭连接，不影响 HTTP API 及其他请求
func WsMiniprogram(c *gin.Context) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("[WS] WsMiniprogram panic recovered: %v", r)
		}
	}()
	conn, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}
	defer conn.Close()

	var userID string
	authOK := false

	// 读取首条消息：auth
	conn.SetReadDeadline(time.Now().Add(15 * time.Second))
	_, msg, err := conn.ReadMessage()
	if err != nil {
		return
	}
	var authMsg struct {
		Type   string `json:"type"`
		UserID string `json:"userId"`
	}
	if json.Unmarshal(msg, &authMsg) == nil && authMsg.Type == "auth" && authMsg.UserID != "" {
		userID = strings.TrimSpace(authMsg.UserID)
		// 占位：校验用户存在即可
		db := database.DB()
		var u model.User
		if db.Where("id = ?", userID).First(&u).Error == nil {
			authOK = true
		}
	}
	if !authOK {
		conn.WriteJSON(map[string]interface{}{"type": "error", "message": "auth failed"})
		return
	}

	// 鉴权通过，开始处理心跳
	conn.SetReadDeadline(time.Time{}) // 取消超时
	client := redis.Client()
	if client == nil {
		log.Printf("[WS] Redis 未启用，在线状态不可用")
		return
	}
	key := wsOnlinePrefix + userID
	ctx := context.Background()
	ttl := time.Duration(wsOfflineTimeout) * time.Second

	// 立即写入一次在线
	client.Set(ctx, key, "1", ttl)

	// 心跳读取循环
	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			break
		}
		var m struct {
			Type string `json:"type"`
		}
		if json.Unmarshal(msg, &m) == nil && (m.Type == "ping" || m.Type == "heartbeat") {
			client.Set(ctx, key, "1", ttl)
			conn.WriteJSON(map[string]interface{}{"type": "pong"})
		}
	}
}

// AdminUsersOnlineStats GET /api/admin/users/online-stats 管理端在线人数统计
// 容错：Redis 不可用时返回 success + onlineCount: 0，不影响管理端其他功能
func AdminUsersOnlineStats(c *gin.Context) {
	client := redis.Client()
	if client == nil {
		c.JSON(http.StatusOK, gin.H{"success": true, "onlineCount": 0})
		return
	}
	ctx := context.Background()
	iter := client.Scan(ctx, 0, wsOnlinePrefix+"*", 0).Iterator()
	count := 0
	for iter.Next(ctx) {
		count++
	}
	if err := iter.Err(); err != nil {
		log.Printf("[WS] AdminUsersOnlineStats Redis scan err: %v，降级返回 0", err)
		c.JSON(http.StatusOK, gin.H{"success": true, "onlineCount": 0})
		return
	}
	c.JSON(http.StatusOK, gin.H{"success": true, "onlineCount": count})
}

// GetWsURL 返回小程序可用的 WSS 地址（基于 API_BASE_URL 派生）
func GetWsURL() string {
	cfg := config.Get()
	if cfg == nil {
		return ""
	}
	base := strings.TrimSuffix(cfg.BaseURL, "/")
	if base == "" {
		return ""
	}
	if strings.HasPrefix(base, "https://") {
		return "wss" + strings.TrimPrefix(base, "https") + "/ws/miniprogram"
	}
	if strings.HasPrefix(base, "http://") {
		return strings.Replace(base, "http", "ws", 1) + "/ws/miniprogram"
	}
	return ""
}