go-ts/internal/client/packet.go

package client

import (
	"bytes"
	"encoding/binary"
	"log"
	"math"
	"time"

	"go-ts/pkg/protocol"
)

func (c *Client) handlePacket(pkt *protocol.Packet) error {
	// log.Printf("Received Packet: ID=%d, Type=%v, Len=%d", pkt.Header.PacketID, pkt.Header.PacketType(), len(pkt.Data))

	switch pkt.Header.PacketType() {
	case protocol.PacketTypeInit1:
		return c.handleInit(pkt)
	case protocol.PacketTypeCommand:
		// Send ACK
		ackData := make([]byte, 2)
		binary.BigEndian.PutUint16(ackData, pkt.Header.PacketID)

		ack := protocol.NewPacket(protocol.PacketTypeAck, ackData)
		// Spec/ts3j: Ack has its own counter
		c.AckPacketID++
		ack.Header.PacketID = c.AckPacketID
		ack.Header.ClientID = c.ClientID
		// ACKs usually don't have NewProtocol flag set in Header byte
		ack.Header.Type &= ^uint8(protocol.PacketFlagNewProtocol)

		// ACKs for Command packets after handshake must be encrypted
		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,
			}
			key, nonce = crypto.GenerateKeyNonce(&ack.Header, true) // Client->Server=true
		}

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

		encData, mac, _ := protocol.EncryptEAX(key, nonce, meta, ack.Data)
		ack.Data = encData
		copy(ack.Header.MAC[:], mac)
		// log.Printf("Sending ACK for server Command PID=%d", pkt.Header.PacketID)

		c.Conn.SendPacket(ack)

		return c.handleCommand(pkt)
	case protocol.PacketTypeCommandLow:
		// Send ACK Low
		ackData := make([]byte, 2)
		binary.BigEndian.PutUint16(ackData, pkt.Header.PacketID)

		ack := protocol.NewPacket(protocol.PacketTypeAckLow, ackData)
		// Spec/ts3j: AckLow has its own counter
		c.AckLowPacketID++
		ack.Header.PacketID = c.AckLowPacketID
		ack.Header.ClientID = c.ClientID
		// ACKs usually don't have NewProtocol flag set in Header byte
		ack.Header.Type &= ^uint8(protocol.PacketFlagNewProtocol)

		// ACKs for Command packets after handshake must be encrypted
		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,
			}
			key, nonce = crypto.GenerateKeyNonce(&ack.Header, true) // Client->Server=true
		}

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

		encData, mac, _ := protocol.EncryptEAX(key, nonce, meta, ack.Data)
		ack.Data = encData
		copy(ack.Header.MAC[:], mac)
		// log.Printf("Sending ACK Low for server CommandLow PID=%d", pkt.Header.PacketID)

		c.Conn.SendPacket(ack)

		return c.handleCommand(pkt)
	case protocol.PacketTypeVoice:
		c.handleVoice(pkt)
	case protocol.PacketTypePing:
		// Respond with Pong
		// Respond with Pong
		pong := protocol.NewPacket(protocol.PacketTypePong, nil)
		// Spec/ts3j: Pong has its own counter
		c.PongPacketID++
		pong.Header.PacketID = c.PongPacketID
		pong.Header.ClientID = c.ClientID

		// Determine valid keys for encryption
		key := protocol.HandshakeKey
		nonce := protocol.HandshakeNonce
		shouldEncrypt := false

		if c.Handshake != nil && c.Handshake.Step >= 6 && len(c.Handshake.SharedMac) > 0 {
			shouldEncrypt = true
			copy(pong.Header.MAC[:], c.Handshake.SharedMac)

			// Generate EAX keys
			if len(c.Handshake.SharedIV) > 0 {
				crypto := &protocol.CryptoState{
					SharedIV:     c.Handshake.SharedIV,
					SharedMac:    c.Handshake.SharedMac,
					GenerationID: 0,
				}
				key, nonce = crypto.GenerateKeyNonce(&pong.Header, true) // Client->Server
			}
		} else {
			// Pre-handshake or fallback
			pong.Header.Type = uint8(protocol.PacketTypePong) | protocol.PacketFlagUnencrypted
			for i := 0; i < 8; i++ {
				pong.Header.MAC[i] = 0
			}
		}

		// The body of the Pong must contain the PID of the Ping it's acknowledging
		pongData := make([]byte, 2)
		binary.BigEndian.PutUint16(pongData, pkt.Header.PacketID)

		if shouldEncrypt {
			// Encrypt the Pong data
			// Meta for Client->Server: PID(2) + CID(2) + PT(1) = 5 bytes
			meta := make([]byte, 5)
			binary.BigEndian.PutUint16(meta[0:2], pong.Header.PacketID)
			binary.BigEndian.PutUint16(meta[2:4], pong.Header.ClientID)
			meta[4] = pong.Header.Type // ensure NewProtocol is NOT set (0x05)

			encData, mac, _ := protocol.EncryptEAX(key, nonce, meta, pongData)
			pong.Data = encData
			copy(pong.Header.MAC[:], mac)
			log.Printf("Sending Encrypted Pong (HeaderPID=%d) for Ping", pong.Header.PacketID)
		} else {
			pong.Data = pongData
			log.Printf("Sending Unencrypted Pong (HeaderPID=%d) for Ping", pong.Header.PacketID)
		}

		c.Conn.SendPacket(pong)
	case protocol.PacketTypePong:
		// Server acknowledged our Ping - calculate RTT
		receivedAt := time.Now()

		// Decrypt Pong body if needed to get the PingID it's acknowledging
		var pongData []byte
		if pkt.Header.FlagUnencrypted() {
			pongData = pkt.Data
		} else 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,
			}
			key, nonce := crypto.GenerateKeyNonce(&pkt.Header, false) // Server->Client
			meta := make([]byte, 3)
			binary.BigEndian.PutUint16(meta[0:2], pkt.Header.PacketID)
			meta[2] = pkt.Header.Type
			decrypted, err := protocol.DecryptEAX(key, nonce, meta, pkt.Data, pkt.Header.MAC[:])
			if err == nil {
				pongData = decrypted
			}
		}

		// Extract PingID from Pong body
		var pingID uint16
		if len(pongData) >= 2 {
			pingID = binary.BigEndian.Uint16(pongData[0:2])
		} else {
			pingID = pkt.Header.PacketID // fallback
		}

		// Calculate RTT if we have the send time
		if sentTime, ok := c.PingSentTimes[pingID]; ok {
			rtt := receivedAt.Sub(sentTime).Seconds() * 1000 // RTT in ms
			delete(c.PingSentTimes, pingID)

			// Update rolling average using Welford's algorithm
			c.PingSampleCount++
			if c.PingSampleCount == 1 {
				c.PingRTT = rtt
				c.PingDeviation = 0
			} else {
				oldMean := c.PingRTT
				c.PingRTT = oldMean + (rtt-oldMean)/float64(c.PingSampleCount)
				// Rolling deviation: exponential smoothing
				c.PingDeviation = c.PingDeviation*0.9 + math.Abs(rtt-c.PingRTT)*0.1
			}

			log.Printf("Received Pong for Ping %d: RTT=%.2fms, AvgRTT=%.2fms, Dev=%.2fms",
				pingID, rtt, c.PingRTT, c.PingDeviation)
		} else {
			log.Printf("Received Pong for unknown Ping %d", pingID)
		}
	case protocol.PacketTypeAck:
		// Server acknowledged our packet
		var data []byte
		var err error

		if pkt.Header.FlagUnencrypted() {
			data = pkt.Data
		} else {
			// ACKs are encrypted
			key := protocol.HandshakeKey
			nonce := protocol.HandshakeNonce

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

			meta := make([]byte, 3) // Server->Client is 3 bytes
			binary.BigEndian.PutUint16(meta[0:2], pkt.Header.PacketID)
			meta[2] = pkt.Header.Type

			data, err = protocol.DecryptEAX(key, nonce, meta, pkt.Data, pkt.Header.MAC[:])
			if err != nil {
				// Try fallback to HandshakeKey if SharedSecret failed
				if !bytes.Equal(key, protocol.HandshakeKey[:]) {
					log.Printf("ACK SharedSecret decrypt failed, trying HandshakeKey...")
					key = protocol.HandshakeKey[:]
					nonce = protocol.HandshakeNonce[:]
					data, err = protocol.DecryptEAX(key, nonce, meta, pkt.Data, pkt.Header.MAC[:])
				}

				if err != nil {
					log.Printf("ACK decryption failed (PID=%d): %v", pkt.Header.PacketID, err)
					return nil
				}
			}
		}

		ackPId := uint16(0)
		if len(data) >= 2 {
			ackPId = binary.BigEndian.Uint16(data[0:2])
		}
		log.Printf("Received ACK for PacketID %d (HeaderPID=%d)", ackPId, pkt.Header.PacketID)

		// If ACK is for clientek (PID=1), proceed with clientinit
		if ackPId == 1 && c.Handshake != nil && c.Handshake.Step == 5 {
			log.Println("clientek acknowledged! Sending clientinit...")
			c.Handshake.Step = 6
			return c.sendClientInit()
		}
		// If ACK is for clientinit (PID=2), we're connected!
		if ackPId == 2 && c.Handshake != nil && c.Handshake.Step == 6 {
			log.Println("clientinit acknowledged! Connection established!")
			c.Connected = true
		}
	}
	return nil
}

func (c *Client) handleInit(pkt *protocol.Packet) error {
	// Determine step based on packet content or local state
	// Simple state machine
	if c.Handshake.Step == 0 {
		if err := c.Handshake.HandlePacket1(pkt); err != nil {
			return err
		}
		log.Println("Handshake Step 1 Completed. Sending Step 2...")
		return c.Handshake.SendPacket2()
	} else if c.Handshake.Step == 1 {
		// Wait, step 1 is processed, we sent step 2.
		// We expect Step 3.
		if pkt.Data[0] == 0x03 {
			if err := c.Handshake.HandlePacket3(pkt); err != nil {
				return err
			}
			log.Println("Handshake Step 3 Completed. Sending Step 4 (Puzzle Solution)...")
			if err := c.Handshake.SendPacket4(); err != nil {
				return err
			}
		}
	}
	return nil
}