DebTik/internal/bwtest/server.go

package bwtest

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

// ServerSession tracks an active incoming bandwidth test session on the server
type serverSession struct {
	id          string
	sessionHash uint64
	clientTCP   net.Conn
	clientAddr  net.IP
	clientPort  int
	handshake   SessionHandshake
	cancel      context.CancelFunc
	ctx         context.Context

	// Statistics
	rxBytes     atomic.Int64
	rxPackets   atomic.Int64
	txBytes     atomic.Int64
	txPackets   atomic.Int64
	lostPackets atomic.Int64
	jitterMs    atomic.Uint64 // float64 bits

	maxSeq     atomic.Uint64
	seqTracker sync.Map // track recently seen sequences for loss detection

	mu     sync.Mutex
	jitter JitterCalculator
}

// BwTestServer represents the local Bandwidth Test server
type BwTestServer struct {
	config ServerConfig
	mu     sync.RWMutex

	tcpListener net.Listener
	udpConn     *net.UDPConn
	running     bool
	cancel      context.CancelFunc

	sessions map[string]*serverSession
	sessMu   sync.RWMutex

	totalServed atomic.Int64
	totalBytes  atomic.Int64
}

// NewServer creates a new BwTestServer instance
func NewServer(cfg ServerConfig) *BwTestServer {
	if cfg.Port <= 0 {
		cfg.Port = 2000
	}
	if cfg.ListenAddr == "" {
		cfg.ListenAddr = "0.0.0.0"
	}
	return &BwTestServer{
		config:   cfg,
		sessions: make(map[string]*serverSession),
	}
}

// Start launches the TCP control listener and UDP data listener
func (s *BwTestServer) Start() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.running {
		return nil
	}

	ctx, cancel := context.WithCancel(context.Background())
	s.cancel = cancel

	tcpAddr := fmt.Sprintf("%s:%d", s.config.ListenAddr, s.config.Port)
	ln, err := net.Listen("tcp", tcpAddr)
	if err != nil {
		cancel()
		return fmt.Errorf("failed to bind BwTest TCP listener on %s: %w", tcpAddr, err)
	}
	s.tcpListener = ln

	udpAddr, err := net.ResolveUDPAddr("udp", tcpAddr)
	if err != nil {
		ln.Close()
		cancel()
		return fmt.Errorf("failed to resolve BwTest UDP address: %w", err)
	}

	uconn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		ln.Close()
		cancel()
		return fmt.Errorf("failed to bind BwTest UDP listener on %s: %w", tcpAddr, err)
	}
	_ = uconn.SetReadBuffer(8 * 1024 * 1024) // 8MB OS UDP socket buffer
	s.udpConn = uconn
	s.running = true

	go s.acceptLoop(ctx)
	go s.udpReceiverLoop(ctx)

	log.Printf("[BwTestServer] Started Bandwidth Test Server listening on %s (TCP & UDP)", tcpAddr)
	return nil
}

// Stop terminates the server and all active sessions
func (s *BwTestServer) Stop() {
	s.mu.Lock()
	if !s.running {
		s.mu.Unlock()
		return
	}
	s.running = false
	if s.cancel != nil {
		s.cancel()
	}
	if s.tcpListener != nil {
		_ = s.tcpListener.Close()
	}
	if s.udpConn != nil {
		_ = s.udpConn.Close()
	}
	s.mu.Unlock()

	s.sessMu.Lock()
	for _, sess := range s.sessions {
		if sess.cancel != nil {
			sess.cancel()
		}
		if sess.clientTCP != nil {
			_ = sess.clientTCP.Close()
		}
	}
	s.sessions = make(map[string]*serverSession)
	s.sessMu.Unlock()

	log.Printf("[BwTestServer] Bandwidth Test Server stopped")
}

// Status returns current operational metrics
func (s *BwTestServer) Status() ServerStatus {
	s.mu.RLock()
	defer s.mu.RUnlock()

	s.sessMu.RLock()
	activeCount := len(s.sessions)
	s.sessMu.RUnlock()

	return ServerStatus{
		Enabled:             s.config.Enabled,
		Port:                s.config.Port,
		Running:             s.running,
		ActiveSessions:      activeCount,
		TotalSessionsServed: s.totalServed.Load(),
		TotalBytesProcessed: s.totalBytes.Load(),
	}
}

// SetPort updates the port and restarts the server if currently running
func (s *BwTestServer) SetPort(port int) error {
	s.mu.Lock()
	wasRunning := s.running
	s.config.Port = port
	s.mu.Unlock()

	if wasRunning {
		s.Stop()
		return s.Start()
	}
	return nil
}

func (s *BwTestServer) acceptLoop(ctx context.Context) {
	for {
		conn, err := s.tcpListener.Accept()
		if err != nil {
			select {
			case <-ctx.Done():
				return
			default:
				time.Sleep(50 * time.Millisecond)
				continue
			}
		}

		go s.handleClientControl(ctx, conn)
	}
}

func (s *BwTestServer) handleClientControl(serverCtx context.Context, conn net.Conn) {
	defer conn.Close()

	remoteTCPAddr, _ := conn.RemoteAddr().(*net.TCPAddr)
	clientIP := net.IPv4(127, 0, 0, 1)
	if remoteTCPAddr != nil {
		clientIP = remoteTCPAddr.IP
	}

	reader := bufio.NewReader(conn)
	line, err := reader.ReadBytes('\n')
	if err != nil {
		return
	}

	var req SessionHandshake
	if err := json.Unmarshal(line, &req); err != nil {
		resp := SessionHandshake{Type: "reject", Message: "Invalid handshake JSON"}
		b, _ := json.Marshal(resp)
		_, _ = conn.Write(append(b, '\n'))
		return
	}

	if req.Type != "request" {
		return
	}

	// Session parameter sanitization
	if req.SessionID == "" {
		req.SessionID = fmt.Sprintf("bw-%d-%d", time.Now().UnixNano(), rand.Intn(10000))
	}
	if req.PacketSize <= 0 {
		req.PacketSize = 1472
	}
	if req.PacketSize < 64 {
		req.PacketSize = 64
	}
	if req.PacketSize > 65507 {
		req.PacketSize = 65507
	}
	if req.Streams <= 0 {
		req.Streams = 1
	}
	if req.Streams > 8 {
		req.Streams = 8
	}
	if req.Duration < 0 {
		req.Duration = 10
	}

	sessCtx, sessCancel := context.WithCancel(serverCtx)
	if req.Duration > 0 {
		sessCtx, sessCancel = context.WithTimeout(serverCtx, time.Duration(req.Duration+2)*time.Second)
	}
	defer sessCancel()

	sess := &serverSession{
		id:          req.SessionID,
		sessionHash: HashSessionID(req.SessionID),
		clientTCP:   conn,
		clientAddr:  clientIP,
		clientPort:  req.UDPPort,
		handshake:   req,
		ctx:         sessCtx,
		cancel:      sessCancel,
	}

	s.sessMu.Lock()
	s.sessions[sess.id] = sess
	s.sessMu.Unlock()

	s.totalServed.Add(1)

	defer func() {
		s.sessMu.Lock()
		delete(s.sessions, sess.id)
		s.sessMu.Unlock()
	}()

	// Respond with accept
	s.mu.RLock()
	listenPort := s.config.Port
	s.mu.RUnlock()

	acceptResp := SessionHandshake{
		Type:      "accept",
		SessionID: sess.id,
		UDPPort:   listenPort,
		Protocol:  req.Protocol,
		Direction: req.Direction,
		Message:   "Ready",
	}
	acceptBytes, _ := json.Marshal(acceptResp)
	if _, err := conn.Write(append(acceptBytes, '\n')); err != nil {
		return
	}

	// Launch worker routines based on direction and protocol
	if req.Protocol == "udp" {
		// If client is receiving (server transmitting)
		if req.Direction == "rx" || req.Direction == "both" {
			go s.runServerUDPTx(sessCtx, sess)
		}
	} else if req.Protocol == "tcp" {
		if req.Direction == "rx" || req.Direction == "both" {
			go s.runServerTCPTx(sessCtx, sess)
		}
		if req.Direction == "tx" || req.Direction == "both" {
			go s.runServerTCPRx(sessCtx, sess)
		}
	}

	// Read control messages / reports from client until termination
	go s.reportLoop(sessCtx, sess)

	for {
		line, err := reader.ReadBytes('\n')
		if err != nil {
			break
		}
		var msg SessionHandshake
		if err := json.Unmarshal(line, &msg); err == nil {
			if msg.Type == "stop" {
				break
			}
		}
	}
}

// udpReceiverLoop processes incoming UDP test packets across all active sessions
func (s *BwTestServer) udpReceiverLoop(ctx context.Context) {
	buf := make([]byte, 65535)

	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		n, remoteAddr, err := s.udpConn.ReadFromUDP(buf)
		if err != nil {
			select {
			case <-ctx.Done():
				return
			default:
				time.Sleep(10 * time.Millisecond)
				continue
			}
		}

		if n < HeaderSize {
			continue
		}

		hdr, err := DecodeHeader(buf[:n])
		if err != nil {
			continue
		}

		// Find matching session
		s.sessMu.RLock()
		var matched *serverSession
		for _, sess := range s.sessions {
			if sess.sessionHash == hdr.SessionHash {
				matched = sess
				break
			}
		}
		s.sessMu.RUnlock()

		if matched == nil {
			continue
		}

		// Update remote client UDP address for reverse transmission
		if matched.clientPort == 0 || matched.clientAddr == nil {
			matched.clientAddr = remoteAddr.IP
			matched.clientPort = remoteAddr.Port
		}

		nowNano := time.Now().UnixNano()
		matched.rxBytes.Add(int64(n))
		matched.rxPackets.Add(1)
		s.totalBytes.Add(int64(n))

		// Jitter calculation
		matched.mu.Lock()
		curJitter := matched.jitter.Update(hdr.TimestampNano, nowNano)
		matched.mu.Unlock()

		// Packet loss tracking via sequence gaps
		lastMax := matched.maxSeq.Load()
		if hdr.Sequence > lastMax {
			if lastMax > 0 && hdr.Sequence > lastMax+1 {
				gap := int64(hdr.Sequence - lastMax - 1)
				matched.lostPackets.Add(gap)
			}
			matched.maxSeq.Store(hdr.Sequence)
		}

		_ = curJitter
	}
}

// runServerUDPTx sends UDP packets from server to client (when client requested RX or Both)
func (s *BwTestServer) runServerUDPTx(ctx context.Context, sess *serverSession) {
	// Wait briefly for client to send first packet so we know client's UDP port, or use negotiated port
	targetPort := sess.clientPort
	if targetPort == 0 {
		for i := 0; i < 20; i++ {
			if sess.clientPort != 0 {
				targetPort = sess.clientPort
				break
			}
			time.Sleep(50 * time.Millisecond)
		}
	}
	if targetPort == 0 {
		targetPort = s.config.Port
	}

	dstUDPAddr := &net.UDPAddr{
		IP:   sess.clientAddr,
		Port: targetPort,
	}

	pktSize := sess.handshake.PacketSize
	if pktSize < HeaderSize {
		pktSize = HeaderSize
	}
	buf := make([]byte, pktSize)

	limiter := NewRateLimiter(sess.handshake.BandwidthLimit)
	var seq uint64

	hdr := PacketHeader{
		Magic:         MagicBytes,
		Version:       ProtoVer,
		PacketType:    PacketTypeData,
		StreamID:      0,
		Flags:         0,
		SessionHash:   sess.sessionHash,
		PayloadLength: uint32(pktSize - HeaderSize),
	}

	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		if limiter != nil {
			limiter.Wait(pktSize)
		}

		seq++
		hdr.Sequence = seq
		hdr.TimestampNano = time.Now().UnixNano()
		EncodeHeader(buf, hdr)

		n, err := s.udpConn.WriteToUDP(buf, dstUDPAddr)
		if err == nil {
			sess.txBytes.Add(int64(n))
			sess.txPackets.Add(1)
			s.totalBytes.Add(int64(n))
		}
	}
}

// runServerTCPTx generates continuous TCP stream data to client
func (s *BwTestServer) runServerTCPTx(ctx context.Context, sess *serverSession) {
	chunk := make([]byte, 64*1024)
	limiter := NewRateLimiter(sess.handshake.BandwidthLimit)

	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		if limiter != nil {
			limiter.Wait(len(chunk))
		}

		n, err := sess.clientTCP.Write(chunk)
		if err != nil {
			return
		}
		sess.txBytes.Add(int64(n))
		s.totalBytes.Add(int64(n))
	}
}

// runServerTCPRx reads data sent from client over TCP
func (s *BwTestServer) runServerTCPRx(ctx context.Context, sess *serverSession) {
	buf := make([]byte, 64*1024)
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		n, err := sess.clientTCP.Read(buf)
		if err != nil {
			return
		}
		sess.rxBytes.Add(int64(n))
		s.totalBytes.Add(int64(n))
	}
}

// reportLoop periodically sends feedback reports back to the client over TCP
func (s *BwTestServer) reportLoop(ctx context.Context, sess *serverSession) {
	ticker := time.NewTicker(300 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			report := SessionHandshake{
				Type:        "report",
				SessionID:   sess.id,
				RxBytes:     sess.rxBytes.Load(),
				RxPackets:   sess.rxPackets.Load(),
				LostPackets: sess.lostPackets.Load(),
			}
			b, err := json.Marshal(report)
			if err == nil {
				_, _ = sess.clientTCP.Write(append(b, '\n'))
			}
		}
	}
}