nnet/internal/connection/websocket_connection.go

package connection

import (
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/rs/xid"
)

// WebSocketConnection WebSocket连接包装器
type WebSocketConnection struct {
	id         string
	conn       *websocket.Conn
	mu         sync.RWMutex
	createdAt  time.Time
	lastActive time.Time
	attributes map[string]interface{}
}

// NewWebSocketConnection 创建WebSocket连接
// 如果id为空，将自动生成唯一的xid
func NewWebSocketConnection(id string, conn *websocket.Conn) *WebSocketConnection {
	if id == "" {
		id = xid.New().String()
	}
	return &WebSocketConnection{
		id:         id,
		conn:       conn,
		createdAt:  time.Now(),
		lastActive: time.Now(),
		attributes: make(map[string]interface{}),
	}
}

// ID 获取连接ID
func (c *WebSocketConnection) ID() string {
	return c.id
}

// RemoteAddr 获取远程地址
func (c *WebSocketConnection) RemoteAddr() string {
	if c.conn != nil {
		return c.conn.RemoteAddr().String()
	}
	return ""
}

// LocalAddr 获取本地地址
func (c *WebSocketConnection) LocalAddr() string {
	if c.conn != nil {
		return c.conn.LocalAddr().String()
	}
	return ""
}

// Write 写入数据
func (c *WebSocketConnection) Write(data []byte) error {
	c.mu.Lock()
	c.lastActive = time.Now()
	c.mu.Unlock()

	if c.conn != nil {
		return c.conn.WriteMessage(websocket.TextMessage, data)
	}
	return nil
}

// Close 关闭连接
func (c *WebSocketConnection) Close() error {
	if c.conn != nil {
		return c.conn.Close()
	}
	return nil
}

// SetAttribute 设置属性
func (c *WebSocketConnection) SetAttribute(key string, value interface{}) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.attributes == nil {
		c.attributes = make(map[string]interface{})
	}
	c.attributes[key] = value
}

// GetAttribute 获取属性
func (c *WebSocketConnection) GetAttribute(key string) interface{} {
	c.mu.RLock()
	defer c.mu.RUnlock()
	if c.attributes == nil {
		return nil
	}
	return c.attributes[key]
}

// LastActive 获取最后活动时间
func (c *WebSocketConnection) LastActive() time.Time {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.lastActive
}

// UpdateActive 更新活动时间
func (c *WebSocketConnection) UpdateActive() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.lastActive = time.Now()
}