diff --git a/handshake/payload.go b/handshake/payload.go index 4567fc0d..7f6f8de9 100644 --- a/handshake/payload.go +++ b/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.