go-ts/internal/client/packet.go

package client

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

	"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: Header PID for ACK matches the packet being acknowledged
		ack.Header.PacketID = pkt.Header.PacketID
		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.PacketTypeVoice:
		c.handleVoice(pkt)
	case protocol.PacketTypePing:
		// Respond with Pong
		pong := protocol.NewPacket(protocol.PacketTypePong, nil)
		// Spec/ts3j: Header PID for Pong matches the Ping ID
		pong.Header.PacketID = pkt.Header.PacketID
		pong.Header.ClientID = c.ClientID
		// Must NOT have NewProtocol (0x20) flag for Pings/Pongs
		pong.Header.Type = uint8(protocol.PacketTypePong) | protocol.PacketFlagUnencrypted

		// Use SharedMac if available, otherwise zeros
		if c.Handshake != nil && len(c.Handshake.SharedMac) > 0 {
			copy(pong.Header.MAC[:], c.Handshake.SharedMac)
		} else {
			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
		pong.Data = make([]byte, 2)
		binary.BigEndian.PutUint16(pong.Data, pkt.Header.PacketID)

		log.Printf("Sending Pong (HeaderPID=%d) for Ping", pong.Header.PacketID)
		c.Conn.SendPacket(pong)
	case protocol.PacketTypePong:
		// Server acknowledged our Ping
		log.Printf("Received Pong for sequence %d", pkt.Header.PacketID)
	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
}