DebTik/internal/bwtest/server.go

package bwtest

import (
	"bufio"
	"context"
	crand "crypto/rand"
	"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
	tcpWriter   *tcpFrameWriter
	clientAddr  net.IP
	clientPort  int
	peerReady   chan struct{}
	peerOnce    sync.Once
	udpPacing   *udpPacer
	rxTracker   udpTracker
	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

	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()
	_ = conn.SetReadDeadline(time.Now().Add(5 * time.Second))

	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
	}

	if req.Protocol == "tcp" && req.DataFraming != tcpFramingVersion {
		b, _ := json.Marshal(SessionHandshake{Type: "reject", Message: "TCP test requires framed protocol; upgrade both peers"})
		_, _ = conn.Write(append(b, '\n'))
		return
	}
	if req.Protocol == "udp" && req.BandwidthLimit <= 0 && req.UDPFeedback != udpFeedbackVersion {
		b, _ := json.Marshal(SessionHandshake{Type: "reject", Message: "Automatic UDP requires receiver feedback; upgrade both peers"})
		_, _ = conn.Write(append(b, '\n'))
		return
	}
	if req.Protocol == "udp" {
		if req.PacketSize <= 0 {
			req.PacketSize = 1472
		}
		size, err := udpPayloadSize(req.PacketSize, clientIP)
		if err != nil {
			return
		}
		req.PacketSize = size
	}

	_ = conn.SetReadDeadline(time.Time{})

	// 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
	}

	var sessCtx context.Context
	var sessCancel context.CancelFunc
	if req.Duration > 0 {
		sessCtx, sessCancel = context.WithTimeout(serverCtx, time.Duration(req.Duration+2)*time.Second)
	} else {
		sessCtx, sessCancel = context.WithCancel(serverCtx)
	}
	defer sessCancel()
	stopClose := context.AfterFunc(sessCtx, func() { conn.Close() })
	defer stopClose()

	sess := &serverSession{
		id:          req.SessionID,
		sessionHash: HashSessionID(req.SessionID),
		clientTCP:   conn,
		tcpWriter:   &tcpFrameWriter{conn: conn},
		clientAddr:  clientIP,
		clientPort:  req.UDPPort,
		peerReady:   make(chan struct{}),
		udpPacing:   newUDPPacer(req.BandwidthLimit),
		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",
	}
	if req.Protocol == "tcp" {
		acceptResp.DataFraming = tcpFramingVersion
	} else {
		acceptResp.UDPFeedback = udpFeedbackVersion
		acceptResp.PacketSize = req.PacketSize
	}
	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)
		}

	}

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

	frameBuf := make([]byte, tcpFrameMax)
	for {
		var line []byte
		if req.Protocol == "tcp" {
			kind, payload, err := readTCPFrame(reader, frameBuf)
			if err != nil {
				break
			}
			if kind == tcpFrameData {
				sess.rxBytes.Add(int64(len(payload)))
				sess.rxPackets.Add(1)
				s.totalBytes.Add(int64(len(payload)))
				continue
			}
			line = payload
		} else {
			var err error
			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
			}
			if msg.Type == "report" && msg.SessionID == sess.id {
				sess.udpPacing.observe(msg.RxBytes, msg.RxPackets, msg.LostPackets, time.Now())
			}
		}
	}

}

// 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
		}

		// Learn the actual source port even for RX-only clients behind NAT.
		// Restrict discovery to the host that owns this TCP control session.
		if !remoteAddr.IP.Equal(matched.clientAddr) {
			continue
		}
		if hdr.Version != ProtoVer {
			continue
		}
		if hdr.PacketType == udpProbeType {
			matched.mu.Lock()
			matched.clientPort = remoteAddr.Port
			matched.mu.Unlock()
			matched.peerOnce.Do(func() { close(matched.peerReady) })
			continue
		}
		if hdr.PacketType != PacketTypeData || int(hdr.PayloadLength) != n-HeaderSize {
			continue
		}
		accepted, lost := matched.rxTracker.accept(hdr.StreamID, hdr.Sequence)
		matched.lostPackets.Store(lost)
		if !accepted {
			continue
		}
		matched.mu.Lock()
		matched.clientPort = remoteAddr.Port
		matched.mu.Unlock()
		matched.peerOnce.Do(func() { close(matched.peerReady) })

		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()

		_ = curJitter
	}
}

// runServerUDPTx sends UDP packets from server to client (when client requested RX or Both)
func (s *BwTestServer) runServerUDPTx(ctx context.Context, sess *serverSession) {
	if sess.handshake.UDPFeedback == udpFeedbackVersion {
		select {
		case <-sess.peerReady:
		case <-ctx.Done():
			return
		}
	}
	sess.mu.Lock()
	port := sess.clientPort
	sess.mu.Unlock()
	if port <= 0 {
		return
	}
	dstUDPAddr := &net.UDPAddr{IP: sess.clientAddr, Port: port}

	pktSize := sess.handshake.PacketSize
	if pktSize < HeaderSize {
		pktSize = HeaderSize
	}
	buf := make([]byte, pktSize)
	if _, err := crand.Read(buf); err != nil {
		return
	}

	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 !sess.udpPacing.wait(ctx, pktSize) {
			return
		}

		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)
	if _, err := crand.Read(chunk); err != nil {
		return
	}
	limiter := NewRateLimiter(sess.handshake.BandwidthLimit)

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

		if limiter != nil {
			if !limiter.WaitContext(ctx, len(chunk)) {
				return
			}
		}

		n, err := sess.tcpWriter.write(tcpFrameData, chunk)
		if err != nil {
			return
		}
		sess.txBytes.Add(int64(n))
		sess.txPackets.Add(1)
		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 {
				if sess.handshake.Protocol == "tcp" {
					_, _ = sess.tcpWriter.write(tcpFrameControl, b)
				} else {
					_, _ = sess.clientTCP.Write(append(b, '\n'))
				}
			}
		}
	}
}