Feature: Add Go TCP server framework

This commit is contained in:
Shaun
2025-12-19 14:12:29 +01:00
committed by yuanyuanxiang
parent 7d2cf647ec
commit 5a33628b92
16 changed files with 2807 additions and 0 deletions

View File

@@ -0,0 +1,197 @@
package connection
import (
"fmt"
"net"
"sync"
"sync/atomic"
"time"
"github.com/yuanyuanxiang/SimpleRemoter/server/go/buffer"
)
// ClientInfo stores client metadata
type ClientInfo struct {
ClientID string // Client ID from login (MasterID)
IP string
ComputerName string
OS string
CPU string
HasCamera bool
Version string
InstallTime string
LoginTime time.Time
ClientType string
FilePath string
GroupName string
}
// Context represents a client connection context
type Context struct {
ID uint64
Conn net.Conn
RemoteAddr string
// Buffers
InBuffer *buffer.Buffer // Received compressed data
OutBuffer *buffer.Buffer // Decompressed data for processing
// Client info
Info ClientInfo
IsLoggedIn atomic.Bool
// Connection state
OnlineTime time.Time
lastActiveNs atomic.Int64 // Unix nanoseconds for thread-safe access
// Protocol state
CompressMethod int
FlagType FlagType
HeaderLen int
FlagLen int
HeaderEncType int // Header encryption type (0-7)
HeaderParams []byte // Header parameters for decoding (flag bytes)
// User data - for storing dialog/handler references
UserData interface{}
// Internal
mu sync.RWMutex
closed atomic.Bool
sendLock sync.Mutex
server *Manager
}
// FlagType represents the protocol flag type
type FlagType int
const (
FlagUnknown FlagType = iota
FlagShine
FlagFuck
FlagHello
FlagHell
FlagWinOS
)
// Compression methods
const (
CompressUnknown = -2
CompressZlib = -1
CompressZstd = 0
CompressNone = 1
)
// NewContext creates a new connection context
func NewContext(conn net.Conn, mgr *Manager) *Context {
now := time.Now()
ctx := &Context{
Conn: conn,
RemoteAddr: conn.RemoteAddr().String(),
InBuffer: buffer.New(8192),
OutBuffer: buffer.New(8192),
OnlineTime: now,
CompressMethod: CompressZstd,
FlagType: FlagUnknown,
server: mgr,
}
ctx.lastActiveNs.Store(now.UnixNano())
return ctx
}
// Send sends data to the client (thread-safe)
func (c *Context) Send(data []byte) error {
if c.closed.Load() {
return ErrConnectionClosed
}
c.sendLock.Lock()
defer c.sendLock.Unlock()
_, err := c.Conn.Write(data)
if err != nil {
return err
}
c.UpdateLastActive()
return nil
}
// UpdateLastActive updates the last active time (thread-safe)
func (c *Context) UpdateLastActive() {
c.lastActiveNs.Store(time.Now().UnixNano())
}
// LastActive returns the last active time (thread-safe)
func (c *Context) LastActive() time.Time {
return time.Unix(0, c.lastActiveNs.Load())
}
// TimeSinceLastActive returns duration since last activity (thread-safe)
func (c *Context) TimeSinceLastActive() time.Duration {
return time.Since(c.LastActive())
}
// Close closes the connection
func (c *Context) Close() error {
if c.closed.Swap(true) {
return nil // Already closed
}
return c.Conn.Close()
}
// IsClosed returns whether the connection is closed
func (c *Context) IsClosed() bool {
return c.closed.Load()
}
// GetPeerIP returns the peer IP address
func (c *Context) GetPeerIP() string {
if host, _, err := net.SplitHostPort(c.RemoteAddr); err == nil {
return host
}
return c.RemoteAddr
}
// AliveTime returns how long the connection has been alive
func (c *Context) AliveTime() time.Duration {
return time.Since(c.OnlineTime)
}
// SetInfo sets the client info
func (c *Context) SetInfo(info ClientInfo) {
c.mu.Lock()
defer c.mu.Unlock()
c.Info = info
}
// GetInfo returns the client info
func (c *Context) GetInfo() ClientInfo {
c.mu.RLock()
defer c.mu.RUnlock()
return c.Info
}
// SetUserData stores user-defined data
func (c *Context) SetUserData(data interface{}) {
c.mu.Lock()
defer c.mu.Unlock()
c.UserData = data
}
// GetUserData retrieves user-defined data
func (c *Context) GetUserData() interface{} {
c.mu.RLock()
defer c.mu.RUnlock()
return c.UserData
}
// GetClientID returns the client ID for logging
// If ClientID is set (from login), returns it; otherwise returns connection ID as fallback
func (c *Context) GetClientID() string {
c.mu.RLock()
defer c.mu.RUnlock()
if c.Info.ClientID != "" {
return c.Info.ClientID
}
return fmt.Sprintf("conn-%d", c.ID)
}

View File

@@ -0,0 +1,23 @@
package connection
import "errors"
var (
// ErrConnectionClosed indicates the connection is closed
ErrConnectionClosed = errors.New("connection closed")
// ErrServerClosed indicates the server is shut down
ErrServerClosed = errors.New("server closed")
// ErrMaxConnections indicates max connections reached
ErrMaxConnections = errors.New("max connections reached")
// ErrInvalidPacket indicates an invalid packet
ErrInvalidPacket = errors.New("invalid packet")
// ErrUnsupportedProtocol indicates unsupported protocol
ErrUnsupportedProtocol = errors.New("unsupported protocol")
// ErrDecompressFailed indicates decompression failure
ErrDecompressFailed = errors.New("decompression failed")
)

View File

@@ -0,0 +1,115 @@
package connection
import (
"sync"
"sync/atomic"
)
// Manager manages all client connections
type Manager struct {
connections sync.Map // map[uint64]*Context
count atomic.Int64
maxConns int
idCounter atomic.Uint64
// Callbacks
onConnect func(*Context)
onDisconnect func(*Context)
onReceive func(*Context, []byte)
}
// NewManager creates a new connection manager
func NewManager(maxConns int) *Manager {
if maxConns <= 0 {
maxConns = 10000
}
return &Manager{
maxConns: maxConns,
}
}
// SetCallbacks sets the callback functions
func (m *Manager) SetCallbacks(onConnect, onDisconnect func(*Context), onReceive func(*Context, []byte)) {
m.onConnect = onConnect
m.onDisconnect = onDisconnect
m.onReceive = onReceive
}
// Add adds a new connection
func (m *Manager) Add(ctx *Context) error {
if int(m.count.Load()) >= m.maxConns {
return ErrMaxConnections
}
ctx.ID = m.idCounter.Add(1)
m.connections.Store(ctx.ID, ctx)
m.count.Add(1)
if m.onConnect != nil {
m.onConnect(ctx)
}
return nil
}
// Remove removes a connection
func (m *Manager) Remove(ctx *Context) {
if _, ok := m.connections.LoadAndDelete(ctx.ID); ok {
m.count.Add(-1)
if m.onDisconnect != nil {
m.onDisconnect(ctx)
}
}
}
// Get retrieves a connection by ID
func (m *Manager) Get(id uint64) *Context {
if v, ok := m.connections.Load(id); ok {
return v.(*Context)
}
return nil
}
// Count returns the current connection count
func (m *Manager) Count() int {
return int(m.count.Load())
}
// Range iterates over all connections
func (m *Manager) Range(fn func(*Context) bool) {
m.connections.Range(func(key, value interface{}) bool {
return fn(value.(*Context))
})
}
// Broadcast sends data to all connections
func (m *Manager) Broadcast(data []byte) {
m.connections.Range(func(key, value interface{}) bool {
ctx := value.(*Context)
if !ctx.IsClosed() {
_ = ctx.Send(data)
}
return true
})
}
// CloseAll closes all connections
func (m *Manager) CloseAll() {
m.connections.Range(func(key, value interface{}) bool {
ctx := value.(*Context)
_ = ctx.Close()
return true
})
}
// OnReceive calls the receive callback
func (m *Manager) OnReceive(ctx *Context, data []byte) {
if m.onReceive != nil {
m.onReceive(ctx, data)
}
}
// UpdateMaxConnections updates the maximum connections limit
func (m *Manager) UpdateMaxConnections(max int) {
m.maxConns = max
}