Protocol/stream.go

package eq2net

import (
	"encoding/binary"
	"fmt"
	"net"
	"sync"
	"time"
)

// StreamState represents the connection state
type StreamState int

const (
	Established StreamState = iota
	WaitClose
	Closing
	Disconnecting
	Closed
)

// StreamType represents different stream types
type StreamType int

const (
	UnknownStream StreamType = iota
	LoginStream
	WorldStream
	ZoneStream
	ChatOrMailStream
	ChatStream
	MailStream
	EQ2Stream
)

// SessionRequest packet structure
type SessionRequest struct {
	UnknownA  uint32
	Session   uint32
	MaxLength uint32
}

// SessionResponse packet structure
type SessionResponse struct {
	Session   uint32
	Key       uint32
	UnknownA  uint8
	Format    uint8
	UnknownB  uint8
	MaxLength uint32
	UnknownD  uint32
}

// ClientSessionStats for tracking connection statistics
type ClientSessionStats struct {
	RequestID         uint16
	LastLocalDelta    uint32
	AverageDelta      uint32
	LowDelta          uint32
	HighDelta         uint32
	LastRemoteDelta   uint32
	PacketsSent       uint64
	PacketsReceived   uint64
}

// ServerSessionStats for server-side statistics
type ServerSessionStats struct {
	RequestID        uint16
	CurrentTime      uint32
	Unknown1         uint32
	ReceivedPackets  uint32
	Unknown2         uint32
	SentPackets      uint32
	Unknown3         uint32
	SentPackets2     uint32
	Unknown4         uint32
	ReceivedPackets2 uint32
}

// Stream represents a connection stream
type Stream struct {
	mu sync.RWMutex

	// Connection info
	remoteAddr    *net.UDPAddr
	remoteIP      net.IP
	remotePort    uint16
	localAddr     *net.UDPAddr
	sessionID     uint32
	key           uint32
	state         StreamState
	streamType    StreamType

	// Buffers
	buffer         [8192]byte
	oversizeBuffer []byte
	oversizeOffset uint32
	oversizeLength uint32
	rogueBuffer    []byte
	rogueBufOffset uint32
	rogueBufSize   uint32

	// Packet tracking
	receivedPackets uint32
	sentPackets     uint32
	nextInSeq       uint16
	nextOutSeq      uint16
	lastAckSent     uint16
	lastAckReceived uint16

	// Timing
	lastReceiveTime time.Time
	lastSendTime    time.Time
	avgRoundTrip    time.Duration
	minRoundTrip    time.Duration
	maxRoundTrip    time.Duration
	retransmitTimeout time.Duration

	// Queues
	incomingQueue   []*ProtocolPacket
	outgoingQueue   []*ProtocolPacket
	sentQueue       map[uint16]*ProtocolPacket
	futurePackets   map[uint16]*ProtocolPacket
	
	// Options
	compressed      bool
	encoded         bool
	maxLength       uint32
	opcodeSize      uint8

	// Callbacks
	onPacket        func(*ApplicationPacket)
	onDisconnect    func()
}

// NewStream creates a new stream
func NewStream(remoteAddr *net.UDPAddr, streamType StreamType) *Stream {
	s := &Stream{
		remoteAddr:        remoteAddr,
		remoteIP:          remoteAddr.IP,
		remotePort:        uint16(remoteAddr.Port),
		streamType:        streamType,
		state:             Established,
		sentQueue:         make(map[uint16]*ProtocolPacket),
		futurePackets:     make(map[uint16]*ProtocolPacket),
		lastReceiveTime:   time.Now(),
		lastSendTime:      time.Now(),
		retransmitTimeout: 500 * time.Millisecond,
		maxLength:         512,
		opcodeSize:        2,
	}

	// Set opcode size based on stream type
	if streamType == LoginStream || streamType == ChatStream || streamType == MailStream {
		s.opcodeSize = 1
	}

	return s
}

// SetSessionInfo sets the session ID and key
func (s *Stream) SetSessionInfo(sessionID uint32, key uint32) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.sessionID = sessionID
	s.key = key
}

// GetSessionID returns the session ID
func (s *Stream) GetSessionID() uint32 {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.sessionID
}

// GetKey returns the session key
func (s *Stream) GetKey() uint32 {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.key
}

// GetState returns the current stream state
func (s *Stream) GetState() StreamState {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.state
}

// SetState sets the stream state
func (s *Stream) SetState(state StreamState) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.state = state
}

// GetRemoteAddr returns the remote address
func (s *Stream) GetRemoteAddr() *net.UDPAddr {
	return s.remoteAddr
}

// Process processes an incoming packet
func (s *Stream) Process(data []byte) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.lastReceiveTime = time.Now()
	s.receivedPackets++

	// Parse the packet
	if len(data) < 2 {
		return fmt.Errorf("packet too small")
	}

	opcode := binary.BigEndian.Uint16(data)
	
	// Handle different packet types
	switch uint8(opcode & 0xFF) {
	case OPSessionRequest:
		return s.handleSessionRequest(data)
	case OPSessionResponse:
		return s.handleSessionResponse(data)
	case OPKeepAlive:
		return s.handleKeepAlive(data)
	case OPPacket:
		return s.handlePacket(data)
	case OPFragment:
		return s.handleFragment(data)
	case OPAck:
		return s.handleAck(data)
	case OPOutOfOrderAck:
		return s.handleOutOfOrderAck(data)
	case OPSessionDisconnect:
		return s.handleDisconnect(data)
	case OPCombined:
		return s.handleCombined(data)
	case OPAppCombined:
		return s.handleAppCombined(data)
	default:
		return fmt.Errorf("unknown opcode: %02x", opcode&0xFF)
	}
}

// handleSessionRequest handles session request packets
func (s *Stream) handleSessionRequest(data []byte) error {
	if len(data) < 14 { // 2 byte opcode + 12 byte SessionRequest
		return fmt.Errorf("session request too small")
	}

	req := &SessionRequest{}
	req.UnknownA = binary.LittleEndian.Uint32(data[2:6])
	req.Session = binary.LittleEndian.Uint32(data[6:10])
	req.MaxLength = binary.LittleEndian.Uint32(data[10:14])

	s.sessionID = req.Session
	s.maxLength = req.MaxLength

	// Send session response
	return s.sendSessionResponse()
}

// handleSessionResponse handles session response packets
func (s *Stream) handleSessionResponse(data []byte) error {
	if len(data) < 21 { // 2 byte opcode + 19 byte SessionResponse
		return fmt.Errorf("session response too small")
	}

	resp := &SessionResponse{}
	resp.Session = binary.LittleEndian.Uint32(data[2:6])
	resp.Key = binary.LittleEndian.Uint32(data[6:10])
	resp.UnknownA = data[10]
	resp.Format = data[11]
	resp.UnknownB = data[12]
	resp.MaxLength = binary.LittleEndian.Uint32(data[13:17])
	resp.UnknownD = binary.LittleEndian.Uint32(data[17:21])

	s.sessionID = resp.Session
	s.key = resp.Key
	s.maxLength = resp.MaxLength
	s.compressed = (resp.Format & 0x01) != 0
	s.encoded = (resp.Format & 0x04) != 0

	return nil
}

// handleKeepAlive handles keep-alive packets
func (s *Stream) handleKeepAlive(data []byte) error {
	// Send keep-alive response
	return s.sendKeepAliveResponse()
}

// handlePacket handles regular data packets
func (s *Stream) handlePacket(data []byte) error {
	if len(data) < 4 {
		return fmt.Errorf("packet too small")
	}

	// Extract sequence number
	sequence := binary.BigEndian.Uint16(data[2:4])
	
	// Validate CRC
	if !ValidateCRC(data, len(data), s.key) {
		return fmt.Errorf("CRC validation failed")
	}

	// Create protocol packet
	packet := NewProtocolPacket(data[4:len(data)-2], -1) // Skip header and CRC
	packet.Sequence = sequence

	// Handle sequence
	if sequence == s.nextInSeq {
		// In order packet
		s.processInOrderPacket(packet)
		s.nextInSeq++
		
		// Check for any future packets that are now in order
		s.processFuturePackets()
		
		// Send ACK
		s.sendAck(sequence)
	} else if sequence > s.nextInSeq {
		// Future packet - store it
		s.futurePackets[sequence] = packet
		// Send out of order ACK
		s.sendOutOfOrderAck(sequence)
	}
	// Ignore past packets (already processed)

	return nil
}

// handleFragment handles fragmented packets
func (s *Stream) handleFragment(data []byte) error {
	if len(data) < 8 {
		return fmt.Errorf("fragment too small")
	}

	sequence := binary.BigEndian.Uint16(data[2:4])
	totalSize := binary.BigEndian.Uint32(data[4:8])
	
	// Validate CRC
	if !ValidateCRC(data, len(data), s.key) {
		return fmt.Errorf("CRC validation failed")
	}

	// Initialize oversize buffer if needed
	if s.oversizeBuffer == nil || s.oversizeLength != totalSize {
		s.oversizeBuffer = make([]byte, totalSize)
		s.oversizeOffset = 0
		s.oversizeLength = totalSize
	}

	// Copy fragment data
	fragmentData := data[8 : len(data)-2] // Skip header and CRC
	copy(s.oversizeBuffer[s.oversizeOffset:], fragmentData)
	s.oversizeOffset += uint32(len(fragmentData))

	// Send ACK
	s.sendAck(sequence)
	s.nextInSeq = sequence + 1

	// Check if we have the complete packet
	if s.oversizeOffset >= s.oversizeLength {
		// Process the complete packet
		packet := NewProtocolPacket(s.oversizeBuffer[:s.oversizeLength], -1)
		s.processInOrderPacket(packet)
		
		// Reset oversize buffer
		s.oversizeBuffer = nil
		s.oversizeOffset = 0
		s.oversizeLength = 0
	}

	return nil
}

// handleAck handles acknowledgment packets
func (s *Stream) handleAck(data []byte) error {
	if len(data) < 4 {
		return fmt.Errorf("ack too small")
	}

	sequence := binary.BigEndian.Uint16(data[2:4])
	
	// Remove acknowledged packet from sent queue
	if packet, ok := s.sentQueue[sequence]; ok {
		packet.Acked = true
		delete(s.sentQueue, sequence)
		
		// Update RTT based on this ACK
		if packet.SentTime > 0 {
			rtt := time.Since(time.Unix(int64(packet.SentTime), 0))
			s.updateRTT(rtt)
		}
	}
	
	s.lastAckReceived = sequence
	
	return nil
}

// handleOutOfOrderAck handles out-of-order acknowledgments
func (s *Stream) handleOutOfOrderAck(data []byte) error {
	if len(data) < 4 {
		return fmt.Errorf("out of order ack too small")
	}

	sequence := binary.BigEndian.Uint16(data[2:4])
	
	// Resend packets between lastAckReceived and sequence
	for seq := s.lastAckReceived + 1; seq < sequence; seq++ {
		if packet, ok := s.sentQueue[seq]; ok {
			s.resendPacket(packet)
		}
	}
	
	return nil
}

// handleDisconnect handles disconnect packets
func (s *Stream) handleDisconnect(data []byte) error {
	s.state = Disconnecting
	
	// Call disconnect callback
	if s.onDisconnect != nil {
		s.onDisconnect()
	}
	
	return nil
}

// handleCombined handles combined packets
func (s *Stream) handleCombined(data []byte) error {
	if len(data) < 3 {
		return fmt.Errorf("combined packet too small")
	}

	offset := 2 // Skip opcode
	
	for offset < len(data)-2 { // Leave room for CRC
		if offset+1 >= len(data) {
			break
		}
		
		// Read sub-packet length
		length := uint16(data[offset])
		offset++
		
		if length == 0xFF {
			// Extended length
			if offset+2 >= len(data) {
				break
			}
			length = binary.BigEndian.Uint16(data[offset:])
			offset += 2
		}
		
		// Extract sub-packet
		if offset+int(length) > len(data) {
			break
		}
		
		subPacket := data[offset : offset+int(length)]
		offset += int(length)
		
		// Process sub-packet
		s.Process(subPacket)
	}
	
	return nil
}

// handleAppCombined handles application-level combined packets
func (s *Stream) handleAppCombined(data []byte) error {
	// Similar to handleCombined but for application packets
	return s.handleCombined(data)
}

// processInOrderPacket processes an in-order packet
func (s *Stream) processInOrderPacket(packet *ProtocolPacket) {
	// Decompress if needed
	if s.compressed && packet.EQ2Compressed {
		decompressed := make([]byte, 8192)
		size, err := Decompress(packet.Buffer, uint32(len(packet.Buffer)), decompressed, 8192)
		if err == nil {
			packet.Buffer = decompressed[:size]
			packet.Size = size
		}
	}

	// Decode if needed
	if s.encoded && packet.PacketEncrypted {
		ChatDecode(packet.Buffer, len(packet.Buffer), int(s.key))
		packet.PacketEncrypted = false
	}

	// Convert to application packet and call callback
	if s.onPacket != nil {
		appPacket := s.makeApplicationPacket(packet)
		if appPacket != nil {
			s.onPacket(appPacket)
		}
	}
}

// processFuturePackets processes any future packets that are now in order
func (s *Stream) processFuturePackets() {
	for {
		if packet, ok := s.futurePackets[s.nextInSeq]; ok {
			s.processInOrderPacket(packet)
			delete(s.futurePackets, s.nextInSeq)
			s.nextInSeq++
		} else {
			break
		}
	}
}

// makeApplicationPacket converts a protocol packet to an application packet
func (s *Stream) makeApplicationPacket(packet *ProtocolPacket) *ApplicationPacket {
	if len(packet.Buffer) < int(s.opcodeSize) {
		return nil
	}

	var opcode uint16
	var data []byte
	
	if s.opcodeSize == 1 {
		opcode = uint16(packet.Buffer[0])
		data = packet.Buffer[1:]
	} else {
		opcode = binary.LittleEndian.Uint16(packet.Buffer)
		data = packet.Buffer[2:]
	}

	appPacket := NewApplicationPacket(EmuOpcode(opcode), data)
	appPacket.AppOpcodeSize = s.opcodeSize
	
	return appPacket
}

// Send queues a packet for sending
func (s *Stream) Send(packet *ApplicationPacket) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Convert to protocol packet
	data := packet.Serialize()
	protocolPacket := NewProtocolPacketWithOp(uint16(OPPacket), data)
	
	// Add to outgoing queue
	s.outgoingQueue = append(s.outgoingQueue, protocolPacket)
	
	return nil
}

// sendSessionResponse sends a session response
func (s *Stream) sendSessionResponse() error {
	resp := &SessionResponse{
		Session:   s.sessionID,
		Key:       s.key,
		Format:    0,
		MaxLength: s.maxLength,
	}

	if s.compressed {
		resp.Format |= 0x01
	}
	if s.encoded {
		resp.Format |= 0x04
	}

	// Build response packet
	data := make([]byte, 21)
	binary.BigEndian.PutUint16(data[0:2], uint16(OPSessionResponse))
	binary.LittleEndian.PutUint32(data[2:6], resp.Session)
	binary.LittleEndian.PutUint32(data[6:10], resp.Key)
	data[10] = resp.UnknownA
	data[11] = resp.Format
	data[12] = resp.UnknownB
	binary.LittleEndian.PutUint32(data[13:17], resp.MaxLength)
	binary.LittleEndian.PutUint32(data[17:21], resp.UnknownD)

	// Session response doesn't have CRC
	return s.sendRaw(data)
}

// sendKeepAliveResponse sends a keep-alive response
func (s *Stream) sendKeepAliveResponse() error {
	data := make([]byte, 2)
	binary.BigEndian.PutUint16(data, uint16(OPKeepAlive))
	return s.sendWithCRC(data)
}

// sendAck sends an acknowledgment
func (s *Stream) sendAck(sequence uint16) error {
	data := make([]byte, 4)
	binary.BigEndian.PutUint16(data[0:2], uint16(OPAck))
	binary.BigEndian.PutUint16(data[2:4], sequence)
	return s.sendWithCRC(data)
}

// sendOutOfOrderAck sends an out-of-order acknowledgment
func (s *Stream) sendOutOfOrderAck(sequence uint16) error {
	data := make([]byte, 4)
	binary.BigEndian.PutUint16(data[0:2], uint16(OPOutOfOrderAck))
	binary.BigEndian.PutUint16(data[2:4], sequence)
	return s.sendWithCRC(data)
}

// sendWithCRC sends data with CRC
func (s *Stream) sendWithCRC(data []byte) error {
	// Calculate CRC
	crc := CalculateCRC16(data, s.key)
	
	// Append CRC
	fullData := make([]byte, len(data)+2)
	copy(fullData, data)
	binary.BigEndian.PutUint16(fullData[len(data):], crc)
	
	return s.sendRaw(fullData)
}

// sendRaw sends raw data (should be called from factory)
func (s *Stream) sendRaw(data []byte) error {
	s.lastSendTime = time.Now()
	s.sentPackets++
	// The actual sending will be done by StreamFactory
	return nil
}

// resendPacket resends a packet
func (s *Stream) resendPacket(packet *ProtocolPacket) {
	packet.AttemptCount++
	packet.SentTime = int32(time.Now().Unix())
	// Re-queue for sending
	s.outgoingQueue = append(s.outgoingQueue, packet)
}

// updateRTT updates the round-trip time statistics
func (s *Stream) updateRTT(rtt time.Duration) {
	if s.minRoundTrip == 0 || rtt < s.minRoundTrip {
		s.minRoundTrip = rtt
	}
	if rtt > s.maxRoundTrip {
		s.maxRoundTrip = rtt
	}
	
	// Update average (simple moving average)
	if s.avgRoundTrip == 0 {
		s.avgRoundTrip = rtt
	} else {
		s.avgRoundTrip = (s.avgRoundTrip*3 + rtt) / 4
	}
	
	// Update retransmit timeout
	s.retransmitTimeout = s.avgRoundTrip * 3
	if s.retransmitTimeout > 5*time.Second {
		s.retransmitTimeout = 5 * time.Second
	}
}

// SetOnPacket sets the packet callback
func (s *Stream) SetOnPacket(callback func(*ApplicationPacket)) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.onPacket = callback
}

// SetOnDisconnect sets the disconnect callback
func (s *Stream) SetOnDisconnect(callback func()) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.onDisconnect = callback
}