handshake: Marshal a payload in one pass, straight into out

handshake/payload.go

@@ -1,8 +1,10 @@
 package handshake
 
 import (
+	"encoding/binary"
 	"errors"
 	"math"
+	"slices"
 
 	"google.golang.org/protobuf/encoding/protowire"
 )
@@ -35,33 +37,39 @@ const (
 // with NebulaHandshake{Details: NebulaHandshakeDetails{...}}.
 // Returns out (which may be nil), with the marshalled Payload appended to it.
 func MarshalPayload(out []byte, p Payload) []byte {
-	var details []byte
-
-	if len(p.Cert) > 0 {
-		details = protowire.AppendTag(details, fieldCert, protowire.BytesType)
-		details = protowire.AppendBytes(details, p.Cert)
-	}
+	// Everything but the certificate fits on the stack, which leaves the
+	// certificate the only variable size and out grown just once.
+	var scalars [4 * (1 + binary.MaxVarintLen64)]byte
+	rest := scalars[:0]
 	if p.InitiatorIndex != 0 {
-		details = protowire.AppendTag(details, fieldInitiatorIndex, protowire.VarintType)
-		details = protowire.AppendVarint(details, uint64(p.InitiatorIndex))
+		rest = protowire.AppendTag(rest, fieldInitiatorIndex, protowire.VarintType)
+		rest = protowire.AppendVarint(rest, uint64(p.InitiatorIndex))
 	}
 	if p.ResponderIndex != 0 {
-		details = protowire.AppendTag(details, fieldResponderIndex, protowire.VarintType)
-		details = protowire.AppendVarint(details, uint64(p.ResponderIndex))
+		rest = protowire.AppendTag(rest, fieldResponderIndex, protowire.VarintType)
+		rest = protowire.AppendVarint(rest, uint64(p.ResponderIndex))
 	}
 	if p.Time != 0 {
-		details = protowire.AppendTag(details, fieldTime, protowire.VarintType)
-		details = protowire.AppendVarint(details, p.Time)
+		rest = protowire.AppendTag(rest, fieldTime, protowire.VarintType)
+		rest = protowire.AppendVarint(rest, p.Time)
 	}
 	if p.CertVersion != 0 {
-		details = protowire.AppendTag(details, fieldCertVersion, protowire.VarintType)
-		details = protowire.AppendVarint(details, uint64(p.CertVersion))
+		rest = protowire.AppendTag(rest, fieldCertVersion, protowire.VarintType)
+		rest = protowire.AppendVarint(rest, uint64(p.CertVersion))
 	}
 
+	detailsLen := len(rest)
+	if len(p.Cert) > 0 {
+		detailsLen += protowire.SizeTag(fieldCert) + protowire.SizeBytes(len(p.Cert))
+	}
+	out = slices.Grow(out, protowire.SizeTag(1)+protowire.SizeBytes(detailsLen))
 	out = protowire.AppendTag(out, 1, protowire.BytesType)
-	out = protowire.AppendBytes(out, details)
-
-	return out
+	out = protowire.AppendVarint(out, uint64(detailsLen))
+	if len(p.Cert) > 0 {
+		out = protowire.AppendTag(out, fieldCert, protowire.BytesType)
+		out = protowire.AppendBytes(out, p.Cert)
+	}
+	return append(out, rest...)
 }
 
 // UnmarshalPayload decodes a protobuf-encoded NebulaHandshake message.