Steel/pkg/net/packets/packet.go

package packets

import (
	"bytes"
	"cimeyclust.com/steel/pkg/utils"
	"encoding/binary"
	"fmt"
	"net"
	"sort"
)

const (
	IDConnectedPing                  byte = 0x00
	IDUnconnectedPing                byte = 0x01
	IDUnconnectedPingOpenConnections byte = 0x02
	IDConnectedPong                  byte = 0x03
	IDDetectLostConnections          byte = 0x04
	IDOpenConnectionRequest1         byte = 0x05
	IDOpenConnectionReply1           byte = 0x06
	IDOpenConnectionRequest2         byte = 0x07
	IDOpenConnectionReply2           byte = 0x08
	IDConnectionRequest              byte = 0x09
	IDConnectionRequestAccepted      byte = 0x10
	IDNewIncomingConnection          byte = 0x13
	IDDisconnectNotification         byte = 0x15

	IDIncompatibleProtocolVersion byte = 0x19

	IDUnconnectedPong byte = 0x1c
)

// unconnectedMessageSequence is a sequence of bytes which is found in every unconnected message sent in
// RakNet.
var unconnectedMessageSequence = [16]byte{0x00, 0xff, 0xff, 0x00, 0xfe, 0xfe, 0xfe, 0xfe, 0xfd, 0xfd, 0xfd, 0xfd, 0x12, 0x34, 0x56, 0x78}

// writeAddr writes a UDP address to the buffer passed.
func writeAddr(buffer *bytes.Buffer, addr net.UDPAddr) {
	var ver byte = 6
	if addr.IP.To4() != nil {
		ver = 4
	}
	if addr.IP == nil {
		addr.IP = make([]byte, 16)
	}
	_ = buffer.WriteByte(ver)
	if ver == 4 {
		ipBytes := addr.IP.To4()

		_ = buffer.WriteByte(^ipBytes[0])
		_ = buffer.WriteByte(^ipBytes[1])
		_ = buffer.WriteByte(^ipBytes[2])
		_ = buffer.WriteByte(^ipBytes[3])
		_ = binary.Write(buffer, binary.BigEndian, uint16(addr.Port))
	} else {
		_ = binary.Write(buffer, binary.LittleEndian, int16(23)) // syscall.AF_INET6 on Windows.
		_ = binary.Write(buffer, binary.BigEndian, uint16(addr.Port))
		// The IPv6 address is enclosed in two 0 integers.
		_ = binary.Write(buffer, binary.BigEndian, int32(0))
		_, _ = buffer.Write(addr.IP.To16())
		_ = binary.Write(buffer, binary.BigEndian, int32(0))
	}
}

// readAddr decodes a RakNet address from the buffer passed. If not successful, an error is returned.
func readAddr(buffer *bytes.Buffer, addr *net.UDPAddr) error {
	ver, err := buffer.ReadByte()
	if err != nil {
		return err
	}
	if ver == 4 {
		ipBytes := make([]byte, 4)
		if _, err := buffer.Read(ipBytes); err != nil {
			return fmt.Errorf("error reading raknet address ipv4 bytes: %v", err)
		}
		// Construct an IPv4 out of the 4 bytes we just read.
		addr.IP = net.IPv4((-ipBytes[0]-1)&0xff, (-ipBytes[1]-1)&0xff, (-ipBytes[2]-1)&0xff, (-ipBytes[3]-1)&0xff)
		var port uint16
		if err := binary.Read(buffer, binary.BigEndian, &port); err != nil {
			return fmt.Errorf("error reading raknet address port: %v", err)
		}
		addr.Port = int(port)
	} else {
		buffer.Next(2)
		var port uint16
		if err := binary.Read(buffer, binary.LittleEndian, &port); err != nil {
			return fmt.Errorf("error reading raknet address port: %v", err)
		}
		addr.Port = int(port)
		buffer.Next(4)
		addr.IP = make([]byte, 16)
		if _, err := buffer.Read(addr.IP); err != nil {
			return fmt.Errorf("error reading raknet address ipv6 bytes: %v", err)
		}
		buffer.Next(4)
	}
	return nil
}

const (
	// BitFlagDatagram is set for every valid datagram. It is used to identify
	// packets that are datagrams.
	BitFlagDatagram = 0x80
	// BitFlagACK is set for every ACK Packet.
	BitFlagACK = 0x40
	// BitFlagNACK is set for every NACK Packet.
	BitFlagNACK = 0x20
	// BitFlagNeedsBAndAS is set for every datagram with Packet data, but is not
	// actually used.
	BitFlagNeedsBAndAS = 0x04
)

// noinspection GoUnusedConst
const (
	// ReliabilityUnreliable means that the Packet sent could arrive out of
	// order, be duplicated, or just not arrive at all. It is usually used for
	// high frequency packets of which the order does not matter.
	// lint:ignore U1000 While this constant is unused, it is here for the sake
	// of having all reliabilities documented.
	ReliabilityUnreliable byte = iota
	// ReliabilityUnreliableSequenced means that the Packet sent could be
	// duplicated or not arrive at all, but ensures that it is always handled in
	// the right order.
	ReliabilityUnreliableSequenced
	// ReliabilityReliable means that the Packet sent could not arrive, or
	// arrive out of order, but ensures that the Packet is not duplicated.
	ReliabilityReliable
	// ReliabilityReliableOrdered means that every Packet sent arrives, arrives
	// in the right order and is not duplicated.
	ReliabilityReliableOrdered
	// ReliabilityReliableSequenced means that the Packet sent could not arrive,
	// but ensures that the Packet will be in the right order and not be
	// duplicated.
	ReliabilityReliableSequenced

	// SplitFlag is set in the header if the Packet was split. If so, the
	// encapsulation contains additional data about the fragment.
	SplitFlag = 0x10
)

// Packet is an encapsulation around every Packet sent after the connection is
// established. It is
type Packet struct {
	Reliability byte

	Content       []byte
	MessageIndex  utils.Uint24
	sequenceIndex utils.Uint24
	OrderIndex    utils.Uint24

	Split      bool
	SplitCount uint32
	SplitIndex uint32
	SplitID    uint16
}

// Write writes the Packet and its Content to the buffer passed.
func (packet *Packet) Write(b *bytes.Buffer) {
	header := packet.Reliability << 5
	if packet.Split {
		header |= SplitFlag
	}
	b.WriteByte(header)
	_ = binary.Write(b, binary.BigEndian, uint16(len(packet.Content))<<3)
	if packet.reliable() {
		utils.WriteUint24(b, packet.MessageIndex)
	}
	if packet.sequenced() {
		utils.WriteUint24(b, packet.sequenceIndex)
	}
	if packet.sequencedOrOrdered() {
		utils.WriteUint24(b, packet.OrderIndex)
		// Order channel, we don't care about this.
		b.WriteByte(0)
	}
	if packet.Split {
		_ = binary.Write(b, binary.BigEndian, packet.SplitCount)
		_ = binary.Write(b, binary.BigEndian, packet.SplitID)
		_ = binary.Write(b, binary.BigEndian, packet.SplitIndex)
	}
	b.Write(packet.Content)
}

// Read reads a Packet and its Content from the buffer passed.
func (packet *Packet) Read(b *bytes.Buffer) error {
	header, err := b.ReadByte()
	if err != nil {
		return fmt.Errorf("error reading Packet header: %v", err)
	}
	packet.Split = (header & SplitFlag) != 0
	packet.Reliability = (header & 224) >> 5
	var packetLength uint16
	if err := binary.Read(b, binary.BigEndian, &packetLength); err != nil {
		return fmt.Errorf("error reading Packet length: %v", err)
	}
	packetLength >>= 3
	if packetLength == 0 {
		return fmt.Errorf("invalid Packet length: cannot be 0")
	}

	if packet.reliable() {
		packet.MessageIndex, err = utils.ReadUint24(b)
		if err != nil {
			return fmt.Errorf("error reading Packet message index: %v", err)
		}
	}

	if packet.sequenced() {
		packet.sequenceIndex, err = utils.ReadUint24(b)
		if err != nil {
			return fmt.Errorf("error reading Packet sequence index: %v", err)
		}
	}

	if packet.sequencedOrOrdered() {
		packet.OrderIndex, err = utils.ReadUint24(b)
		if err != nil {
			return fmt.Errorf("error reading Packet order index: %v", err)
		}
		// Order channel (byte), we don't care about this.
		b.Next(1)
	}

	if packet.Split {
		if err := binary.Read(b, binary.BigEndian, &packet.SplitCount); err != nil {
			return fmt.Errorf("error reading Packet split count: %v", err)
		}
		if err := binary.Read(b, binary.BigEndian, &packet.SplitID); err != nil {
			return fmt.Errorf("error reading Packet split ID: %v", err)
		}
		if err := binary.Read(b, binary.BigEndian, &packet.SplitIndex); err != nil {
			return fmt.Errorf("error reading Packet split index: %v", err)
		}
	}

	packet.Content = make([]byte, packetLength)
	if n, err := b.Read(packet.Content); err != nil || n != int(packetLength) {
		return fmt.Errorf("not enough data in Packet: %v bytes read but need %v", n, packetLength)
	}
	return nil
}

func (packet *Packet) reliable() bool {
	switch packet.Reliability {
	case ReliabilityReliable,
		ReliabilityReliableOrdered,
		ReliabilityReliableSequenced:
		return true
	}
	return false
}

func (packet *Packet) sequencedOrOrdered() bool {
	switch packet.Reliability {
	case ReliabilityUnreliableSequenced,
		ReliabilityReliableOrdered,
		ReliabilityReliableSequenced:
		return true
	}
	return false
}

func (packet *Packet) sequenced() bool {
	switch packet.Reliability {
	case ReliabilityUnreliableSequenced,
		ReliabilityReliableSequenced:
		return true
	}
	return false
}

const (
	// packetRange indicates a range of packets, followed by the first and the
	// last Packet in the range.
	packetRange = iota
	// packetSingle indicates a single Packet, followed by its sequence number.
	packetSingle
)

// Acknowledgement is an Acknowledgement Packet that may either be an ACK or a
// NACK, depending on the purpose that it is sent with.
type Acknowledgement struct {
	Packets []utils.Uint24
}

// Write encodes an Acknowledgement Packet and returns an error if not
// successful.
func (ack *Acknowledgement) Write(b *bytes.Buffer, mtu uint16) (n int, err error) {
	packets := ack.Packets
	if len(packets) == 0 {
		return 0, binary.Write(b, binary.BigEndian, int16(0))
	}
	buffer := bytes.NewBuffer(nil)
	// Sort packets before encoding to ensure packets are encoded correctly.
	sort.Slice(packets, func(i, j int) bool {
		return packets[i] < packets[j]
	})

	var firstPacketInRange utils.Uint24
	var lastPacketInRange utils.Uint24
	var recordCount int16

	for index, packet := range packets {
		if buffer.Len() >= int(mtu-10) {
			// We must make sure the final Packet length doesn't exceed the MTU
			// size.
			break
		}
		n++
		if index == 0 {
			// The first Packet, set the first and last Packet to it.
			firstPacketInRange = packet
			lastPacketInRange = packet
			continue
		}
		if packet == lastPacketInRange+1 {
			// Packet is still part of the current range, as it's sequenced
			// properly with the last Packet. Set the last Packet in range to
			// the Packet and continue to the next Packet.
			lastPacketInRange = packet
			continue
		} else {
			// We got to the end of a range/single Packet. We need to write
			// those down now.
			if firstPacketInRange == lastPacketInRange {
				// First Packet equals last Packet, so we have a single Packet
				// record. Write down the Packet, and set the first and last
				// Packet to the current Packet.
				if err := buffer.WriteByte(packetSingle); err != nil {
					return 0, err
				}
				utils.WriteUint24(buffer, firstPacketInRange)

				firstPacketInRange = packet
				lastPacketInRange = packet
			} else {
				// There's a gap between the first and last Packet, so we have a
				// range of packets. Write the first and last Packet of the
				// range and set both to the current Packet.
				if err := buffer.WriteByte(packetRange); err != nil {
					return 0, err
				}
				utils.WriteUint24(buffer, firstPacketInRange)
				utils.WriteUint24(buffer, lastPacketInRange)

				firstPacketInRange = packet
				lastPacketInRange = packet
			}
			// Keep track of the amount of records as we need to write that
			// first.
			recordCount++
		}
	}

	// Make sure the last single Packet/range is written, as we always need to
	// know one Packet ahead to know how we should write the current.
	if firstPacketInRange == lastPacketInRange {
		if err := buffer.WriteByte(packetSingle); err != nil {
			return 0, err
		}
		utils.WriteUint24(buffer, firstPacketInRange)
	} else {
		if err := buffer.WriteByte(packetRange); err != nil {
			return 0, err
		}
		utils.WriteUint24(buffer, firstPacketInRange)
		utils.WriteUint24(buffer, lastPacketInRange)
	}
	recordCount++
	if err := binary.Write(b, binary.BigEndian, recordCount); err != nil {
		return 0, err
	}
	if _, err := b.Write(buffer.Bytes()); err != nil {
		return 0, err
	}
	return n, nil
}

// Read decodes an Acknowledgement Packet and returns an error if not
// successful.
func (ack *Acknowledgement) Read(b *bytes.Buffer) error {
	const maxAcknowledgementPackets = 8192
	var recordCount int16
	if err := binary.Read(b, binary.BigEndian, &recordCount); err != nil {
		return err
	}
	for i := int16(0); i < recordCount; i++ {
		recordType, err := b.ReadByte()
		if err != nil {
			return err
		}
		switch recordType {
		case packetRange:
			start, err := utils.ReadUint24(b)
			if err != nil {
				return err
			}
			end, err := utils.ReadUint24(b)
			if err != nil {
				return err
			}
			for pack := start; pack <= end; pack++ {
				ack.Packets = append(ack.Packets, pack)
				if len(ack.Packets) > maxAcknowledgementPackets {
					return fmt.Errorf("maximum amount of packets in acknowledgement exceeded")
				}
			}
		case packetSingle:
			packet, err := utils.ReadUint24(b)
			if err != nil {
				return err
			}
			ack.Packets = append(ack.Packets, packet)
			if len(ack.Packets) > maxAcknowledgementPackets {
				return fmt.Errorf("maximum amount of packets in acknowledgement exceeded")
			}
		}
	}
	return nil
}