go-ts/internal/client/client.go

package client

import (
	"encoding/binary"
	"fmt"
	"log"
	"sync"
	"time"

	"go-ts/pkg/protocol"
	"go-ts/pkg/transport"

	"gopkg.in/hraban/opus.v2"
)

type Channel struct {
	ID       uint64
	ParentID uint64
	Name     string
	Order    uint64
}

// EventHandler is called when events occur
type EventHandler func(eventType string, data map[string]any)

type Client struct {
	Conn      *transport.TS3Conn
	Handshake *HandshakeState

	Nickname string
	ClientID uint16

	// Counters
	PacketIDCounterC2S uint16 // Commands (Type 0x02)
	VoicePacketID      uint16 // Voice (Type 0x00)
	PingPacketID       uint16 // Type 0x04
	PongPacketID       uint16 // Type 0x05
	AckPacketID        uint16 // Type 0x06
	AckLowPacketID     uint16 // Type 0x07

	// Ping RTT tracking
	PingSentTimes   map[uint16]time.Time // Map PingPacketID -> Time sent
	PingRTT         float64              // Rolling average RTT in ms
	PingDeviation   float64              // Rolling deviation in ms
	PingSampleCount int                  // Number of samples for rolling avg

	// State
	Connected  bool
	ServerName string

	// Fragment reassembly (packet queue like ts3j)
	CommandQueue       map[uint16]*protocol.Packet // Packets waiting for reassembly (Type 0x02)
	ExpectedCommandPID uint16                      // Next expected packet ID (Type 0x02)
	FragmentState      bool                        // Toggle: true = collecting, false = ready

	CommandLowQueue       map[uint16]*protocol.Packet // Packets waiting for reassembly (Type 0x03)
	ExpectedCommandLowPID uint16                      // Next expected packet ID (Type 0x03)
	FragmentStateLow      bool                        // Toggle: true = collecting, false = ready

	// Server Data
	Channels map[uint64]*Channel

	// Audio
	VoiceDecoders  map[uint16]*opus.Decoder // Map VID (sender ID) to decoder
	VoiceEncoder   *opus.Encoder            // Encoder for outgoing audio
	VoiceEncoderMu sync.Mutex               // Protects VoiceEncoder

	// Event handler for public API
	eventHandler EventHandler

	// Done channel to signal shutdown
	done chan struct{}
}

func NewClient(nickname string) *Client {
	return &Client{
		Nickname:              nickname,
		PacketIDCounterC2S:    1,
		AckLowPacketID:        1,
		VoicePacketID:         1,
		Channels:              make(map[uint64]*Channel),
		VoiceDecoders:         make(map[uint16]*opus.Decoder),
		CommandQueue:          make(map[uint16]*protocol.Packet),
		ExpectedCommandPID:    0,
		CommandLowQueue:       make(map[uint16]*protocol.Packet),
		ExpectedCommandLowPID: 0,
		PingSentTimes:         make(map[uint16]time.Time),
		PingRTT:               0,
		PingDeviation:         0,
		done:                  make(chan struct{}),
	}
}

// SetEventHandler sets the callback for events
func (c *Client) SetEventHandler(handler EventHandler) {
	c.eventHandler = handler
}

// emitEvent sends an event to the handler if set
func (c *Client) emitEvent(eventType string, data map[string]any) {
	if c.eventHandler != nil {
		c.eventHandler(eventType, data)
	}
}

// Stop signals the client to stop its loops
func (c *Client) Stop() {
	select {
	case <-c.done:
		// Already closed
	default:
		close(c.done)
	}
}

func (c *Client) Connect(address string) error {
	conn, err := transport.NewTS3Conn(address)
	if err != nil {
		return err
	}
	c.Conn = conn
	log.Printf("Connected to UDP. Starting Handshake...")

	// Initialize Handshake State
	hs, err := NewHandshakeState(c.Conn)
	if err != nil {
		return err
	}
	c.Handshake = hs

	// Improve Identity Security Level to 8 (Standard Requirement)
	c.Handshake.ImproveSecurityLevel(8)

	// Send Init1
	if err := c.Handshake.SendPacket0(); err != nil {
		return err
	}

	// Listen Loop
	pktChan := c.Conn.PacketChan()
	ticker := time.NewTicker(3 * time.Second)
	defer ticker.Stop()

	for {
		// Recovery from panics in the main loop
		func() {
			defer func() {
				if r := recover(); r != nil {
					log.Printf("PANIC in Client loop: %v", r)
				}
			}()

			select {
			case <-c.done:
				log.Println("Client loop stopped")
				return
			case pkt := <-pktChan:
				if pkt == nil {
					// Channel closed
					return
				}
				if err := c.handlePacket(pkt); err != nil {
					log.Printf("Error handling packet: %v", err)
				}
			case <-ticker.C:
				if !c.Connected {
					return // Don't send pings if not connected yet
				}
				// Send KeepAlive Ping (Encrypted, No NewProtocol)
				if err := c.sendPing(); err != nil {
					log.Printf("Error sending Ping: %v", err)
				}
			}
		}()
		// Check if we should exit after the inner function
		select {
		case <-c.done:
			return nil
		default:
		}
	}
}

// sendPing sends an encrypted Ping packet WITHOUT the NewProtocol flag
func (c *Client) sendPing() error {
	pType := protocol.PacketTypePing
	pkt := protocol.NewPacket(pType, nil)
	c.PingPacketID++
	pkt.Header.PacketID = c.PingPacketID
	pkt.Header.ClientID = c.ClientID
	// Note: We do NOT set PacketFlagNewProtocol for Pings

	// Encryption
	key := protocol.HandshakeKey[:]
	nonce := protocol.HandshakeNonce[:]

	if c.Handshake != nil && c.Handshake.Step >= 6 && len(c.Handshake.SharedIV) > 0 {
		crypto := &protocol.CryptoState{
			SharedIV:     c.Handshake.SharedIV,
			SharedMac:    c.Handshake.SharedMac,
			GenerationID: 0,
		}
		keyArr, nonceArr := crypto.GenerateKeyNonce(&pkt.Header, true) // Client->Server=true
		key = keyArr
		nonce = nonceArr
	}

	// Meta for Client->Server: PID(2) + CID(2) + PT(1) = 5 bytes
	meta := make([]byte, 5)
	binary.BigEndian.PutUint16(meta[0:2], pkt.Header.PacketID)
	binary.BigEndian.PutUint16(meta[2:4], pkt.Header.ClientID)
	meta[4] = pkt.Header.Type

	encData, mac, err := protocol.EncryptEAX(key, nonce, meta, pkt.Data)
	if err != nil {
		return fmt.Errorf("encryption failed: %w", err)
	}

	pkt.Data = encData
	copy(pkt.Header.MAC[:], mac)

	// Record send time for RTT calculation
	c.PingSentTimes[pkt.Header.PacketID] = time.Now()

	log.Printf("Sending proper Encrypted Ping (PID=%d)", pkt.Header.PacketID)
	return c.Conn.SendPacket(pkt)
}