nnet/internal/server/websocket_server.go

package server

import (
	"github.com/noahlann/nnet/pkg/config"
)

// WebSocketServer WebSocket服务器（基于gnetServer，支持ws和wss）
type WebSocketServer struct {
	*Server
	isWSS bool // 是否为WSS（WebSocket Secure）
}

// NewWebSocketServer 创建WebSocket服务器（ws或wss）
func NewWebSocketServer(cfg *config.Config) (*WebSocketServer, error) {
	if cfg == nil {
		cfg = config.DefaultConfig()
	}

	// 检测是否为WSS
	isWSS := false
	addr := cfg.Addr
	if len(addr) >= 6 && addr[:6] == "wss://" {
		isWSS = true
		// 转换为TCP地址格式
		cfg.Addr = "tcp://" + addr[6:]
		// 启用TLS
		if !cfg.TLSEnabled {
			cfg.TLSEnabled = true
		}
	} else if len(addr) >= 5 && addr[:5] == "ws://" {
		// 转换为TCP地址格式
		cfg.Addr = "tcp://" + addr[5:]
	}

	// 创建TCP服务器（WebSocket基于TCP）
	server, err := NewServer(cfg)
	if err != nil {
		return nil, err
	}

	return &WebSocketServer{
		Server: server,
		isWSS:  isWSS,
	}, nil
}

// NewWSSServer 创建WSS服务器（WebSocket Secure）
func NewWSSServer(cfg *config.Config) (*WebSocketServer, error) {
	if cfg == nil {
		cfg = config.DefaultConfig()
	}

	// 确保TLS已启用
	if !cfg.TLSEnabled {
		cfg.TLSEnabled = true
	}

	// 转换地址格式
	addr := cfg.Addr
	if len(addr) >= 5 && addr[:5] == "ws://" {
		// 将ws://转换为wss://
		cfg.Addr = "wss://" + addr[5:]
	} else if len(addr) < 6 || addr[:6] != "wss://" {
		// 如果没有wss://前缀，添加
		if len(addr) > 0 && addr[0] != ':' {
			cfg.Addr = "wss://" + addr
		} else {
			cfg.Addr = "wss://" + addr
		}
	}

	return NewWebSocketServer(cfg)
}