DebTik/internal/bwtest/client.go

package bwtest

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

var (
	ErrTestAlreadyRunning = errors.New("a bandwidth test is already in progress")
	ErrServerRejected     = errors.New("remote server rejected bandwidth test request")
)

// BwTestClient manages executing bandwidth tests against a remote BwTest server
type BwTestClient struct {
	mu            sync.RWMutex
	running       bool
	currentCancel context.CancelFunc

	status  TestStatus
	samples []*LiveSample

	subscribers               map[chan LiveSample]struct{}
	sampleCount               uint64
	jitterSum, peakTx, peakRx float64
}

// NewClient creates a new BwTestClient instance
func NewClient() *BwTestClient {
	return &BwTestClient{
		subscribers: make(map[chan LiveSample]struct{}),
	}
}

// GetStatus returns the current or most recent test state
func (c *BwTestClient) GetStatus() TestStatus {
	c.mu.RLock()
	defer c.mu.RUnlock()

	st := c.status
	if len(c.samples) > 0 {
		st.CurrentSample = c.samples[len(c.samples)-1]
		st.RecentSamples = append([]*LiveSample(nil), c.samples...)
	}
	return st
}

// Subscribe gives each observer its own bounded queue. Slow observers never
// block the data path; they can resynchronize using GetStatus.
func (c *BwTestClient) Subscribe() (<-chan LiveSample, func()) {
	c.mu.Lock()
	ch := make(chan LiveSample, 16)
	c.subscribers[ch] = struct{}{}
	c.mu.Unlock()
	var once sync.Once
	return ch, func() { once.Do(func() { c.mu.Lock(); delete(c.subscribers, ch); close(ch); c.mu.Unlock() }) }
}

func (c *BwTestClient) recordSampleLocked(sample *LiveSample) {
	c.sampleCount++
	c.jitterSum += sample.JitterMs
	c.peakTx = max(c.peakTx, sample.TxMbps)
	c.peakRx = max(c.peakRx, sample.RxMbps)
	c.samples = append(c.samples, sample)
	if len(c.samples) > 200 {
		c.samples = c.samples[len(c.samples)-200:]
	}
	c.status.ElapsedSec = sample.ElapsedSec
	c.status.CurrentSample = sample
	for ch := range c.subscribers {
		select {
		case ch <- *sample:
		default:
		}
	}
}

// Stop terminates any currently active bandwidth test
func (c *BwTestClient) Stop() {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.running && c.currentCancel != nil {
		c.currentCancel()
	}
}

// Run executes a bandwidth test asynchronously
func (c *BwTestClient) Run(ctx context.Context, req SessionRequest) error {
	c.mu.Lock()
	if c.running {
		c.mu.Unlock()
		return ErrTestAlreadyRunning
	}

	if req.Target == "" {
		req.Target = "127.0.0.1"
	}
	if req.Port <= 0 {
		req.Port = 2000
	}
	if req.Protocol != "tcp" {
		req.Protocol = "udp"
	}
	if req.Direction != "tx" && req.Direction != "rx" && req.Direction != "both" {
		req.Direction = "both"
	}
	if req.PacketSize <= 0 {
		req.PacketSize = 1472
	}
	if req.PacketSize < 64 || req.PacketSize > 65507 {
		c.mu.Unlock()
		return fmt.Errorf("packet size must be between 64 and 65507 bytes")
	}
	if req.Streams <= 0 {
		req.Streams = 1
	}
	if req.Streams > 8 {
		req.Streams = 8
	}
	if req.Duration < 0 {
		req.Duration = 10
	}

	testCtx, cancel := context.WithCancel(context.Background())
	c.currentCancel = cancel
	c.running = true
	c.samples = make([]*LiveSample, 0, 200)
	c.sampleCount, c.jitterSum, c.peakTx, c.peakRx = 0, 0, 0, 0

	sessionID := fmt.Sprintf("bw-%d-%d", time.Now().Unix(), rand.Intn(10000))
	c.status = TestStatus{
		Running:        true,
		SessionID:      sessionID,
		Target:         req.Target,
		Port:           req.Port,
		Protocol:       req.Protocol,
		Direction:      req.Direction,
		BandwidthLimit: req.BandwidthLimit,
		PacketSize:     req.PacketSize,
		Streams:        req.Streams,
		Duration:       req.Duration,
		StartTime:      time.Now(),
	}
	c.mu.Unlock()

	go c.executeTest(testCtx, cancel, req, sessionID)
	return nil
}

func (c *BwTestClient) executeTest(ctx context.Context, cancel context.CancelFunc, req SessionRequest, sessionID string) {
	defer func() {
		cancel()
		c.mu.Lock()
		c.running = false
		c.status.Running = false
		c.finalizeSummary()
		c.mu.Unlock()
	}()

	targetAddr := fmt.Sprintf("%s:%d", req.Target, req.Port)
	dialer := net.Dialer{Timeout: 5 * time.Second}
	controlConn, err := dialer.DialContext(ctx, "tcp", targetAddr)
	if err != nil {
		c.mu.Lock()
		c.status.Error = fmt.Sprintf("Failed to connect to control port %s: %v", targetAddr, err)
		c.mu.Unlock()
		return
	}
	defer controlConn.Close()
	stopClose := context.AfterFunc(ctx, func() { controlConn.Close() })
	defer stopClose()
	_ = controlConn.SetDeadline(time.Now().Add(5 * time.Second))

	// If UDP, bind a local UDP socket
	var udpConn *net.UDPConn
	var clientUDPPort int
	if req.Protocol == "udp" {
		local := controlConn.LocalAddr().(*net.TCPAddr)
		uaddr := &net.UDPAddr{IP: local.IP, Zone: local.Zone}
		var err error
		if err == nil {
			udpConn, err = net.ListenUDP("udp", uaddr)
		}
		if err != nil {
			c.mu.Lock()
			c.status.Error = fmt.Sprintf("Failed to create local UDP socket: %v", err)
			c.mu.Unlock()
			return
		}
		defer udpConn.Close()
		_ = udpConn.SetReadBuffer(8 * 1024 * 1024)
		if la, ok := udpConn.LocalAddr().(*net.UDPAddr); ok {
			clientUDPPort = la.Port
		}
	}

	if req.Protocol == "udp" {
		size, err := udpPayloadSize(req.PacketSize, controlConn.RemoteAddr().(*net.TCPAddr).IP)
		if err != nil {
			c.mu.Lock()
			c.status.Error = err.Error()
			c.mu.Unlock()
			return
		}
		req.PacketSize = size
	}
	// Send handshake request
	hs := SessionHandshake{
		Type:           "request",
		SessionID:      sessionID,
		Protocol:       req.Protocol,
		Direction:      req.Direction,
		BandwidthLimit: req.BandwidthLimit,
		PacketSize:     req.PacketSize,
		Streams:        req.Streams,
		Duration:       req.Duration,
		UDPPort:        clientUDPPort,
	}
	if req.Protocol == "tcp" {
		hs.DataFraming = tcpFramingVersion
	} else {
		hs.UDPFeedback = udpFeedbackVersion
	}
	b, _ := json.Marshal(hs)
	if _, err := controlConn.Write(append(b, '\n')); err != nil {
		c.mu.Lock()
		c.status.Error = fmt.Sprintf("Handshake send failed: %v", err)
		c.mu.Unlock()
		return
	}

	// Read accept response
	reader := bufio.NewReader(controlConn)
	line, err := reader.ReadBytes('\n')
	if err != nil {
		c.mu.Lock()
		c.status.Error = fmt.Sprintf("Server response error: %v", err)
		c.mu.Unlock()
		return
	}

	var acceptResp SessionHandshake
	if err := json.Unmarshal(line, &acceptResp); err != nil || acceptResp.Type != "accept" {
		c.mu.Lock()
		c.status.Error = fmt.Sprintf("Server rejected test: %s", acceptResp.Message)
		c.mu.Unlock()
		return
	}

	if req.Protocol == "tcp" && acceptResp.DataFraming != tcpFramingVersion {
		c.mu.Lock()
		c.status.Error = "Server does not support framed TCP bandwidth tests; upgrade both peers"
		c.mu.Unlock()
		return
	}
	if req.Protocol == "udp" {
		if req.BandwidthLimit <= 0 && acceptResp.UDPFeedback != udpFeedbackVersion {
			c.mu.Lock()
			c.status.Error = "Automatic UDP tests require receiver-feedback support on both peers"
			c.mu.Unlock()
			return
		}
		if acceptResp.PacketSize >= 64 {
			req.PacketSize = min(req.PacketSize, acceptResp.PacketSize)
		}
		c.mu.Lock()
		c.status.PacketSize = req.PacketSize
		c.mu.Unlock()
	}
	_ = controlConn.SetDeadline(time.Time{})
	tcpWriter := &tcpFrameWriter{conn: controlConn}

	serverUDPPort := acceptResp.UDPPort
	if serverUDPPort <= 0 {
		serverUDPPort = req.Port
	}

	// Statistics counters
	var (
		localTxBytes      atomic.Int64
		localTxPackets    atomic.Int64
		localRxBytes      atomic.Int64
		localRxPackets    atomic.Int64
		localLostPackets  atomic.Int64
		serverRxBytes     atomic.Int64
		serverRxPackets   atomic.Int64
		serverLostPackets atomic.Int64

		curJitterMs atomic.Uint64 // float64 in bits
		jitterMu    sync.Mutex
		jitterCalc  JitterCalculator
		rxTracker   udpTracker
	)

	sessionHash := HashSessionID(sessionID)
	udpPacing := newUDPPacer(req.BandwidthLimit)
	var serverUDPAddr *net.UDPAddr

	// One reader owns the buffered connection, including bytes prefetched
	// during the handshake. It dispatches data frames and feedback reports.
	go func() {
		buf := make([]byte, tcpFrameMax)
		for {
			var line []byte
			if req.Protocol == "tcp" {
				kind, payload, err := readTCPFrame(reader, buf)
				if err != nil {
					return
				}
				if kind == tcpFrameData {
					localRxBytes.Add(int64(len(payload)))
					localRxPackets.Add(1)
					continue
				}
				line = payload
			} else {
				var err error
				line, err = reader.ReadBytes('\n')
				if err != nil {
					return
				}
			}
			var rep SessionHandshake
			if err := json.Unmarshal(line, &rep); err == nil && rep.Type == "report" {
				serverRxBytes.Store(rep.RxBytes)
				serverRxPackets.Store(rep.RxPackets)
				serverLostPackets.Store(rep.LostPackets)
				udpPacing.observe(rep.RxBytes, time.Now())
			}
		}
	}()

	var wg sync.WaitGroup

	// UDP Data path
	if req.Protocol == "udp" {
		serverUDPAddr, err = net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", req.Target, serverUDPPort))
		if err != nil {
			c.mu.Lock()
			c.status.Error = fmt.Sprintf("Failed to resolve server UDP address: %v", err)
			c.mu.Unlock()
			return
		}

		// RX-only also sends a probe, so the server learns the real UDP tuple.
		if _, err := udpConn.WriteToUDP(udpProbe(sessionHash), serverUDPAddr); err != nil {
			c.mu.Lock()
			c.status.Error = err.Error()
			c.mu.Unlock()
			return
		}
		// Transmitter (TX or Both)
		if req.Direction == "tx" || req.Direction == "both" {
			for s := 0; s < req.Streams; s++ {
				wg.Add(1)
				streamID := uint8(s)
				go func() {
					defer wg.Done()

					buf := make([]byte, req.PacketSize)
					if _, err := crand.Read(buf); err != nil {
						return
					}
					var seq uint64

					hdr := PacketHeader{
						Magic:         MagicBytes,
						Version:       ProtoVer,
						PacketType:    PacketTypeData,
						StreamID:      streamID,
						SessionHash:   sessionHash,
						PayloadLength: uint32(req.PacketSize - HeaderSize),
					}

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

						if !udpPacing.wait(ctx, req.PacketSize) {
							return
						}

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

						n, err := udpConn.WriteToUDP(buf, serverUDPAddr)
						if err == nil {
							localTxBytes.Add(int64(n))
							localTxPackets.Add(1)
						}
					}
				}()
			}
		}

		// Receiver (RX or Both)
		if req.Direction == "rx" || req.Direction == "both" {
			wg.Add(1)
			go func() {
				defer wg.Done()
				rxBuf := make([]byte, 65535)

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

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

					if n < HeaderSize {
						continue
					}

					hdr, err := DecodeHeader(rxBuf[:n])
					if err != nil || hdr.SessionHash != sessionHash || hdr.Version != ProtoVer || hdr.PacketType != PacketTypeData || int(hdr.PayloadLength) != n-HeaderSize {
						continue
					}
					accepted, lost := rxTracker.accept(hdr.StreamID, hdr.Sequence)
					localLostPackets.Store(lost)
					if !accepted {
						continue
					}

					nowNano := time.Now().UnixNano()
					localRxBytes.Add(int64(n))
					localRxPackets.Add(1)

					// Jitter calculation
					jitterMu.Lock()
					j := jitterCalc.Update(hdr.TimestampNano, nowNano)
					jitterMu.Unlock()
					curJitterMs.Store(uint64(j * 1000)) // ms * 1000

				}
			}()
		}
	} else if req.Protocol == "tcp" {
		// TCP Data path
		if req.Direction == "tx" || req.Direction == "both" {
			wg.Add(1)
			go func() {
				defer wg.Done()
				chunk := make([]byte, 64*1024)
				if _, err := crand.Read(chunk); err != nil {
					return
				}
				limiter := NewRateLimiter(req.BandwidthLimit)

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

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

					n, err := tcpWriter.write(tcpFrameData, chunk)
					if err != nil {
						return
					}
					localTxBytes.Add(int64(n))
					localTxPackets.Add(1)
				}
			}()
		}

	}

	// Telemetry sampling loop (every 500ms)
	sampleTicker := time.NewTicker(500 * time.Millisecond)
	defer sampleTicker.Stop()

	var lastTxBytes, lastRxBytes, lastOfferedBytes int64
	lastTime := time.Now()
	startTime := lastTime

	timeoutChan := make(<-chan time.Time)
	if req.Duration > 0 {
		timer := time.NewTimer(time.Duration(req.Duration) * time.Second)
		defer timer.Stop()
		timeoutChan = timer.C
	}

	for {
		select {
		case <-ctx.Done():
			if req.Protocol == "tcp" {
				return
			}
			// Notify remote server to stop
			stopMsg := SessionHandshake{Type: "stop", SessionID: sessionID}
			sb, _ := json.Marshal(stopMsg)
			_, _ = controlConn.Write(append(sb, '\n'))
			return

		case <-timeoutChan:
			if req.Protocol == "tcp" {
				return
			}
			stopMsg := SessionHandshake{Type: "stop", SessionID: sessionID}
			sb, _ := json.Marshal(stopMsg)
			_, _ = controlConn.Write(append(sb, '\n'))
			return

		case now := <-sampleTicker.C:
			intervalSec := now.Sub(lastTime).Seconds()
			if intervalSec <= 0 {
				intervalSec = 0.5
			}
			lastTime = now
			elapsedSec := now.Sub(startTime).Seconds()

			var curTxBytes, curRxBytes int64
			var curLostPackets int64

			if req.Direction != "rx" {
				curTxBytes = serverRxBytes.Load()
				curLostPackets += serverLostPackets.Load()
			}
			if req.Direction != "tx" {
				curRxBytes = localRxBytes.Load()
				curLostPackets += localLostPackets.Load()
			}
			offered := localTxBytes.Load()
			offeredMbps := float64((offered-lastOfferedBytes)*8) / intervalSec / 1e6
			lastOfferedBytes = offered
			if req.Protocol == "udp" && req.Direction != "tx" {
				rep := SessionHandshake{Type: "report", SessionID: sessionID, RxBytes: localRxBytes.Load(), RxPackets: localRxPackets.Load(), LostPackets: localLostPackets.Load()}
				data, _ := json.Marshal(rep)
				_ = controlConn.SetWriteDeadline(time.Now().Add(2 * time.Second))
				if _, err := controlConn.Write(append(data, '\n')); err != nil {
					c.mu.Lock()
					c.status.Error = fmt.Sprintf("Receiver feedback failed: %v", err)
					c.mu.Unlock()
					return
				}
				// Retry discovery while no reverse packet has arrived.
				if localRxPackets.Load() == 0 {
					_, _ = udpConn.WriteToUDP(udpProbe(sessionHash), serverUDPAddr)
				}
			}

			deltaTx := curTxBytes - lastTxBytes
			deltaRx := curRxBytes - lastRxBytes
			lastTxBytes = curTxBytes
			lastRxBytes = curRxBytes

			txBps := int64(float64(deltaTx*8) / intervalSec)
			rxBps := int64(float64(deltaRx*8) / intervalSec)

			txMbps := float64(txBps) / 1e6
			rxMbps := float64(rxBps) / 1e6

			// Total packets & loss %
			totalExpectedPkts := int64(0)
			if req.Direction != "tx" {
				totalExpectedPkts += localRxPackets.Load() + localLostPackets.Load()
			}
			if req.Direction != "rx" {
				totalExpectedPkts += serverRxPackets.Load() + serverLostPackets.Load()
			}
			var lossPct float64
			if totalExpectedPkts > 0 {
				lossPct = (float64(curLostPackets) / float64(totalExpectedPkts)) * 100.0
				if lossPct > 100.0 {
					lossPct = 100.0
				}
			}

			jitter := float64(curJitterMs.Load()) / 1000.0

			sample := &LiveSample{
				PacketSize:    req.PacketSize,
				OfferedTxMbps: offeredMbps, OfferedTxBytes: offered,
				Timestamp:   now,
				ElapsedSec:  elapsedSec,
				TxBps:       txBps,
				RxBps:       rxBps,
				TxMbps:      txMbps,
				RxMbps:      rxMbps,
				TxPackets:   serverRxPackets.Load(),
				RxPackets:   localRxPackets.Load(),
				TxBytes:     curTxBytes,
				RxBytes:     curRxBytes,
				LostPackets: curLostPackets,
				LossPct:     lossPct,
				JitterMs:    jitter,
			}

			c.mu.Lock()
			c.recordSampleLocked(sample)
			c.mu.Unlock()
		}
	}
}

func (c *BwTestClient) finalizeSummary() {
	if len(c.samples) == 0 {
		return
	}

	last := c.samples[len(c.samples)-1]
	duration := last.ElapsedSec
	var avgTx, avgRx, avgLoss, avgJitter float64
	if duration > 0 {
		avgTx = float64(last.TxBytes) * 8 / duration / 1e6
		avgRx = float64(last.RxBytes) * 8 / duration / 1e6
	}
	expected := last.TxPackets + last.RxPackets + last.LostPackets
	if expected > 0 {
		avgLoss = float64(last.LostPackets) * 100 / float64(expected)
	}
	if c.sampleCount > 0 {
		avgJitter = c.jitterSum / float64(c.sampleCount)
	}

	c.status.Summary = &TestSummary{
		DurationSec:      duration,
		StartTime:        c.status.StartTime,
		EndTime:          time.Now(),
		TotalTxBytes:     last.TxBytes,
		TotalRxBytes:     last.RxBytes,
		AvgTxMbps:        avgTx,
		AvgRxMbps:        avgRx,
		PeakTxMbps:       c.peakTx,
		PeakRxMbps:       c.peakRx,
		TotalTxPackets:   last.TxPackets,
		TotalRxPackets:   last.RxPackets,
		TotalLostPackets: last.LostPackets,
		AvgLossPct:       avgLoss,
		AvgJitterMs:      avgJitter,
	}
}