From a8f9c4bee589e44bf7a95cfa2a2bba24f449bc3e Mon Sep 17 00:00:00 2001
From: Richard Grover <38992201+richgrov@users.noreply.github.com>
Date: Fri, 31 Mar 2023 19:57:16 -0600
Subject: [PATCH] Write entire compressed packet at once (#236)

---
 net/packet/packet.go | 110 ++++++++++++-------------------------------
 net/packet/types.go  |  30 ++++++++----
 2 files changed, 52 insertions(+), 88 deletions(-)

diff --git a/net/packet/packet.go b/net/packet/packet.go
index 280fa8e..a0ab2bf 100644
--- a/net/packet/packet.go
+++ b/net/packet/packet.go
@@ -63,103 +63,53 @@ func (p *Packet) packWithoutCompression(w io.Writer) error {
 	VarInt(p.ID).WriteTo(buffer)
 	buffer.Write(p.Data)
 
-	payloadLen := uint32(buffer.Len() - 3)
+	// Write length at front
+	payloadLen := VarInt(buffer.Len() - 3)
+	varIntOffset := 3 - payloadLen.Len()
+	payloadLen.WriteToBytes(buffer.Bytes()[varIntOffset:])
 
-	// Determine where to start writing the header based on the payload size
-	var headerStart int
-	if payloadLen <= 0xFF>>1 {
-		headerStart = 2
-	} else if payloadLen <= 0xFFFF>>2 {
-		headerStart = 1
-	} else if payloadLen <= 0xFFFFFF>>3 {
-		headerStart = 0
-	} else {
-		panic(fmt.Errorf("packet length %d is too large", payloadLen))
-	}
-
-	// Write the packet length at the beginning of the packet
-	for i := headerStart; payloadLen != 0; i++ {
-		b := byte(payloadLen & 0b01111111)
-		payloadLen >>= 7
-
-		if payloadLen != 0 {
-			b |= 0b10000000
-		}
-
-		buffer.Bytes()[i] = b
-	}
-	_, err := w.Write(buffer.Bytes()[headerStart:])
+	_, err := w.Write(buffer.Bytes()[varIntOffset:])
 	return err
 }
 
 func (p *Packet) packWithCompression(w io.Writer, threshold int) error {
 	buff := bufPool.Get().(*bytes.Buffer)
 	defer bufPool.Put(buff)
+	// Allocate room for the 'packet length' and 'data length' fields. Each can take up to 3 bytes
 	buff.Reset()
+	buff.Write([]byte{0, 0, 0, 0, 0, 0})
+
+	var writeStart int
 
 	if len(p.Data) < threshold {
-		_, err := VarInt(0).WriteTo(buff)
-		if err != nil {
-			return err
-		}
-		_, err = VarInt(p.ID).WriteTo(buff)
-		if err != nil {
-			return err
-		}
-		_, err = buff.Write(p.Data)
-		if err != nil {
-			return err
-		}
-		// Packet Length
-		_, err = VarInt(buff.Len()).WriteTo(w)
-		if err != nil {
-			return err
-		}
-		// Data Length + Packet ID + Data
-		_, err = buff.WriteTo(w)
-		if err != nil {
-			return err
-		}
+		VarInt(p.ID).WriteTo(buff)
+		buff.Write(p.Data)
+		// Packet is below compression threshold so 'data length' is 0
+		// Front of the packet is already initialized to 0, so just decrement the offset
+		writeStart = 5
 	} else {
 		zw := zlib.NewWriter(buff)
-		n1, err := VarInt(p.ID).WriteTo(zw)
-		if err != nil {
-			return err
-		}
-		n2, err := zw.Write(p.Data)
-		if err != nil {
-			return err
-		}
-		err = zw.Close()
+		varIntLen, _ := VarInt(p.ID).WriteTo(zw)
+		zw.Write(p.Data)
+
+		err := zw.Close()
 		if err != nil {
 			return err
 		}
 
-		dataLength := bufPool.Get().(*bytes.Buffer)
-		defer bufPool.Put(dataLength)
-		dataLength.Reset()
-		n3, err := VarInt(int(n1) + n2).WriteTo(dataLength)
-		if err != nil {
-			return err
-		}
-
-		// Packet Length
-		_, err = VarInt(int(n3) + buff.Len()).WriteTo(w)
-		if err != nil {
-			return err
-		}
-		// Data Length
-		_, err = dataLength.WriteTo(w)
-		if err != nil {
-			return err
-		}
-		// PacketID + Data
-		_, err = buff.WriteTo(w)
-		if err != nil {
-			return err
-		}
+		// Write 'data length' before ID + payload
+		uncompressedLen := VarInt(varIntLen + int64(len(p.Data)))
+		writeStart = 6 - uncompressedLen.Len()
+		uncompressedLen.WriteToBytes(buff.Bytes()[writeStart:])
 	}
-	return nil
+
+	// Write 'packet length' before all other fields
+	packetLen := VarInt(buff.Len() - writeStart)
+	start := writeStart - packetLen.Len()
+	VarInt(packetLen).WriteToBytes(buff.Bytes()[start:])
+
+	_, err := w.Write(buff.Bytes()[start:])
+	return err
 }
 
 // UnPack in-place decompression a packet
diff --git a/net/packet/types.go b/net/packet/types.go
index 050fc84..c58e690 100644
--- a/net/packet/types.go
+++ b/net/packet/types.go
@@ -252,21 +252,28 @@ func (l *Long) ReadFrom(r io.Reader) (n int64, err error) {
 }
 
 func (v VarInt) WriteTo(w io.Writer) (n int64, err error) {
-	vi := make([]byte, 0, MaxVarIntLen)
+	vi := make([]byte, MaxVarIntLen)
+	nn := v.WriteToBytes(vi)
+	_, err = w.Write(vi[:nn])
+	return int64(nn), err
+}
+
+func (v VarInt) WriteToBytes(buf []byte) int {
 	num := uint32(v)
+	i := 0
 	for {
 		b := num & 0x7F
 		num >>= 7
 		if num != 0 {
 			b |= 0x80
 		}
-		vi = append(vi, byte(b))
+		buf[i] = byte(b)
+		i++
 		if num == 0 {
 			break
 		}
 	}
-	nn, err := w.Write(vi)
-	return int64(nn), err
+	return i
 }
 
 func (v *VarInt) ReadFrom(r io.Reader) (n int64, err error) {
@@ -308,21 +315,28 @@ func (v VarInt) Len() int {
 }
 
 func (v VarLong) WriteTo(w io.Writer) (n int64, err error) {
-	vi := make([]byte, 0, MaxVarLongLen)
+	vi := make([]byte, MaxVarLongLen)
+	nn := v.WriteToBytes(vi)
+	_, err = w.Write(vi[:nn])
+	return int64(nn), err
+}
+
+func (v VarLong) WriteToBytes(buf []byte) int {
 	num := uint64(v)
+	i := 0
 	for {
 		b := num & 0x7F
 		num >>= 7
 		if num != 0 {
 			b |= 0x80
 		}
-		vi = append(vi, byte(b))
+		buf[i] = byte(b)
+		i++
 		if num == 0 {
 			break
 		}
 	}
-	nn, err := w.Write(vi)
-	return int64(nn), err
+	return i
 }
 
 func (v *VarLong) ReadFrom(r io.Reader) (n int64, err error) {