From 65dd82ad9075059b62a945fb3121f55495a0c4e0 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sun, 28 Aug 2022 22:51:03 +0300
Subject: [PATCH] implement a quicvarint.Append method

---
 quicvarint/varint.go      | 19 +++++++++++++++++++
 quicvarint/varint_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/quicvarint/varint.go b/quicvarint/varint.go
index 66e0d39e..28dcaa93 100644
--- a/quicvarint/varint.go
+++ b/quicvarint/varint.go
@@ -88,6 +88,25 @@ func Write(w Writer, i uint64) {
 	}
 }
 
+func Append(b []byte, i uint64) []byte {
+	if i <= maxVarInt1 {
+		return append(b, uint8(i))
+	}
+	if i <= maxVarInt2 {
+		return append(b, []byte{uint8(i>>8) | 0x40, uint8(i)}...)
+	}
+	if i <= maxVarInt4 {
+		return append(b, []byte{uint8(i>>24) | 0x80, uint8(i >> 16), uint8(i >> 8), uint8(i)}...)
+	}
+	if i <= maxVarInt8 {
+		return append(b, []byte{
+			uint8(i>>56) | 0xc0, uint8(i >> 48), uint8(i >> 40), uint8(i >> 32),
+			uint8(i >> 24), uint8(i >> 16), uint8(i >> 8), uint8(i),
+		}...)
+	}
+	panic(fmt.Sprintf("%#x doesn't fit into 62 bits", i))
+}
+
 // WriteWithLen writes i in the QUIC varint format with the desired length to w.
 func WriteWithLen(w Writer, i uint64, length protocol.ByteCount) {
 	if length != 1 && length != 2 && length != 4 && length != 8 {
diff --git a/quicvarint/varint_test.go b/quicvarint/varint_test.go
index acf4a31c..71db3a79 100644
--- a/quicvarint/varint_test.go
+++ b/quicvarint/varint_test.go
@@ -2,6 +2,7 @@ package quicvarint
 
 import (
 	"bytes"
+	"math/rand"
 
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
@@ -191,6 +192,34 @@ var _ = Describe("Varint encoding / decoding", func() {
 				Expect(Read(b)).To(BeEquivalentTo(494878333))
 			})
 		})
+
+		Context("appending", func() {
+			It("panics when given a too large number (> 62 bit)", func() {
+				Expect(func() { Append(nil, maxVarInt8+1) }).Should(Panic())
+			})
+
+			It("appends", func() {
+				for i := 0; i < 10000; i++ {
+					var limit int64
+					switch rand.Int() % 4 {
+					case 0:
+						limit = maxVarInt1
+					case 1:
+						limit = maxVarInt2
+					case 2:
+						limit = maxVarInt4
+					case 3:
+						limit = maxVarInt8
+					}
+
+					n := uint64(rand.Int63n(limit))
+					b := Append(nil, n)
+					buf := &bytes.Buffer{}
+					Write(buf, n)
+					Expect(b).To(Equal(buf.Bytes()))
+				}
+			})
+		})
 	})
 
 	Context("determining the length needed for encoding", func() {