eq2go/internal/udp/connection.go

package udp

import (
	"crypto/rand"
	"encoding/binary"
	"net"
	"sync"
	"time"
)

type ConnectionState int

const (
	StateClosed ConnectionState = iota
	StateEstablished
	StateClosing
)

type Connection struct {
	addr    *net.UDPAddr
	conn    *net.UDPConn
	handler PacketHandler
	state   ConnectionState
	mutex   sync.RWMutex

	// Session data
	sessionID  uint32
	key        uint32
	compressed bool
	encoded    bool
	maxLength  uint32

	// Sequence tracking
	nextInSeq  uint16
	nextOutSeq uint16

	// Queues
	inboundQueue  []*ApplicationPacket
	outboundQueue []*ProtocolPacket
	ackQueue      []uint16

	// Timing
	lastPacketTime time.Time

	// Crypto
	crypto *Crypto
}

func NewConnection(addr *net.UDPAddr, conn *net.UDPConn, handler PacketHandler) *Connection {
	return &Connection{
		addr:           addr,
		conn:           conn,
		handler:        handler,
		state:          StateClosed,
		maxLength:      512,
		lastPacketTime: time.Now(),
		crypto:         NewCrypto(),
	}
}

func (c *Connection) ProcessPacket(data []byte) {
	c.lastPacketTime = time.Now()

	packet, err := ParseProtocolPacket(data)
	if err != nil {
		return
	}

	switch packet.Opcode {
	case OpSessionRequest:
		c.handleSessionRequest(packet)
	case OpSessionResponse:
		c.handleSessionResponse(packet)
	case OpPacket:
		c.handleDataPacket(packet)
	case OpAck:
		c.handleAck(packet)
	case OpKeepAlive:
		c.sendKeepAlive()
	case OpSessionDisconnect:
		c.Close()
	}
}

func (c *Connection) handleSessionRequest(packet *ProtocolPacket) {
	if len(packet.Data) < 12 {
		return
	}

	// Parse session request
	c.sessionID = binary.LittleEndian.Uint32(packet.Data[4:8])
	requestedMaxLen := binary.LittleEndian.Uint32(packet.Data[8:12])

	if requestedMaxLen > 0 {
		c.maxLength = requestedMaxLen
	}

	// Generate encryption key
	keyBytes := make([]byte, 4)
	rand.Read(keyBytes)
	c.key = binary.LittleEndian.Uint32(keyBytes)

	// Send session response
	c.sendSessionResponse()
	c.state = StateEstablished
}

func (c *Connection) handleSessionResponse(packet *ProtocolPacket) {
	// Client-side session response handling
	if len(packet.Data) < 20 {
		return
	}

	c.sessionID = binary.LittleEndian.Uint32(packet.Data[0:4])
	c.key = binary.LittleEndian.Uint32(packet.Data[4:8])
	format := packet.Data[9]
	c.compressed = (format & 0x01) != 0
	c.encoded = (format & 0x04) != 0
	c.maxLength = binary.LittleEndian.Uint32(packet.Data[12:16])

	c.state = StateEstablished
}

func (c *Connection) handleDataPacket(packet *ProtocolPacket) {
	if len(packet.Data) < 2 {
		return
	}

	seq := binary.BigEndian.Uint16(packet.Data[0:2])
	payload := packet.Data[2:]

	// Simple in-order processing for now
	if seq == c.nextInSeq {
		c.nextInSeq++
		c.sendAck(seq)

		// Process application packet
		if appPacket, err := c.processApplicationData(payload); err == nil {
			c.handler.HandlePacket(c, appPacket)
		}
	}
}

func (c *Connection) handleAck(packet *ProtocolPacket) {
	if len(packet.Data) < 2 {
		return
	}

	seq := binary.BigEndian.Uint16(packet.Data[0:2])
	// Remove acknowledged packets from retransmit queue
	_ = seq // TODO: implement retransmit queue
}

func (c *Connection) processApplicationData(data []byte) (*ApplicationPacket, error) {
	// Decrypt if needed
	if c.crypto.IsEncrypted() {
		data = c.crypto.Decrypt(data)
	}

	// Decompress if needed
	if c.compressed && len(data) > 0 {
		var err error
		data, err = Decompress(data)
		if err != nil {
			return nil, err
		}
	}

	return ParseApplicationPacket(data)
}

func (c *Connection) SendPacket(packet *ApplicationPacket) {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	data := packet.Serialize()

	// Compress if needed
	if c.compressed && len(data) > 128 {
		if compressed, err := Compress(data); err == nil {
			data = compressed
		}
	}

	// Encrypt if needed
	if c.crypto.IsEncrypted() {
		data = c.crypto.Encrypt(data)
	}

	// Create protocol packet
	protocolData := make([]byte, 2+len(data))
	binary.BigEndian.PutUint16(protocolData[0:2], c.nextOutSeq)
	copy(protocolData[2:], data)
	c.nextOutSeq++

	protocolPacket := &ProtocolPacket{
		Opcode: OpPacket,
		Data:   protocolData,
	}

	c.sendProtocolPacket(protocolPacket)
}

func (c *Connection) sendSessionResponse() {
	data := make([]byte, 20)
	binary.LittleEndian.PutUint32(data[0:4], c.sessionID)
	binary.LittleEndian.PutUint32(data[4:8], c.key)
	data[8] = 2 // UnknownA

	var format uint8
	if c.compressed {
		format |= 0x01
	}
	if c.encoded {
		format |= 0x04
	}
	data[9] = format

	data[10] = 0 // UnknownB
	binary.LittleEndian.PutUint32(data[12:16], c.maxLength)
	binary.LittleEndian.PutUint32(data[16:20], 0) // UnknownD

	packet := &ProtocolPacket{
		Opcode: OpSessionResponse,
		Data:   data,
	}

	c.sendProtocolPacket(packet)
}

func (c *Connection) sendAck(seq uint16) {
	data := make([]byte, 2)
	binary.BigEndian.PutUint16(data, seq)

	packet := &ProtocolPacket{
		Opcode: OpAck,
		Data:   data,
	}

	c.sendProtocolPacket(packet)
}

func (c *Connection) sendKeepAlive() {
	packet := &ProtocolPacket{
		Opcode: OpKeepAlive,
		Data:   []byte{},
	}

	c.sendProtocolPacket(packet)
}

func (c *Connection) sendProtocolPacket(packet *ProtocolPacket) {
	data := packet.Serialize()
	c.conn.WriteToUDP(data, c.addr)
}

func (c *Connection) Close() {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	if c.state == StateEstablished {
		c.state = StateClosing

		// Send disconnect
		disconnectData := make([]byte, 6)
		binary.LittleEndian.PutUint32(disconnectData[0:4], c.sessionID)
		disconnectData[4] = 0
		disconnectData[5] = 6

		packet := &ProtocolPacket{
			Opcode: OpSessionDisconnect,
			Data:   disconnectData,
		}

		c.sendProtocolPacket(packet)
	}

	c.state = StateClosed
}