67 lines
1.6 KiB
Go
67 lines
1.6 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"net-tunnel/pkg/proto"
|
|
)
|
|
|
|
// ControlConnection 表示客户端和服务端之间的控制连接
|
|
type ControlConnection struct {
|
|
stream proto.TunnelService_EstablishControlConnectionServer
|
|
cancel context.CancelFunc
|
|
ctx context.Context
|
|
}
|
|
|
|
// NewControlConnection 创建新的控制连接
|
|
func NewControlConnection(stream proto.TunnelService_EstablishControlConnectionServer) *ControlConnection {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
return &ControlConnection{
|
|
stream: stream,
|
|
ctx: ctx,
|
|
cancel: cancel,
|
|
}
|
|
}
|
|
|
|
// Send 发送消息到客户端
|
|
func (c *ControlConnection) Send(msg *proto.Message) error {
|
|
return c.stream.Send(msg)
|
|
}
|
|
|
|
// Close 关闭连接
|
|
func (c *ControlConnection) Close() {
|
|
c.cancel()
|
|
}
|
|
|
|
// ConnectionManager 管理所有客户端的控制连接
|
|
type ConnectionManager struct {
|
|
connections sync.Map // map[clientID]*ControlConnection
|
|
}
|
|
|
|
// NewConnectionManager 创建新的连接管理器
|
|
func NewConnectionManager() *ConnectionManager {
|
|
return &ConnectionManager{
|
|
connections: sync.Map{},
|
|
}
|
|
}
|
|
|
|
// AddConnection 添加新的控制连接
|
|
func (m *ConnectionManager) AddConnection(id string, conn *ControlConnection) {
|
|
m.connections.Store(id, conn)
|
|
}
|
|
|
|
// RemoveConnection 移除控制连接
|
|
func (m *ConnectionManager) RemoveConnection(id string) {
|
|
m.connections.Delete(id)
|
|
}
|
|
|
|
// GetConnection 获取指定客户端的控制连接
|
|
func (m *ConnectionManager) GetConnection(id string) (*ControlConnection, bool) {
|
|
conn, ok := m.connections.Load(id)
|
|
if !ok {
|
|
return nil, false
|
|
}
|
|
return conn.(*ControlConnection), true
|
|
}
|