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

This commit is contained in:
JackDoan committed 2026-10-15 10:42:10 -05:00
1 parent 4d42378537
commit 7c2dcd7ba1
1 file changed
+25 -17
+25 -17
View File
@@ -1,8 +1,10 @@
package handshake package handshake
import ( import (
"encoding/binary"
"errors" "errors"
"math" "math"
"slices"
"google.golang.org/protobuf/encoding/protowire" "google.golang.org/protobuf/encoding/protowire"
) )
@@ -35,33 +37,39 @@ const (
// with NebulaHandshake{Details: NebulaHandshakeDetails{...}}. // with NebulaHandshake{Details: NebulaHandshakeDetails{...}}.
// Returns out (which may be nil), with the marshalled Payload appended to it. // Returns out (which may be nil), with the marshalled Payload appended to it.
func MarshalPayload(out []byte, p Payload) []byte { func MarshalPayload(out []byte, p Payload) []byte {
var details []byte // Everything but the certificate fits on the stack, which leaves the
// certificate the only variable size and out grown just once.
if len(p.Cert) > 0 { var scalars [4 * (1 + binary.MaxVarintLen64)]byte
details = protowire.AppendTag(details, fieldCert, protowire.BytesType) rest := scalars[:0]
details = protowire.AppendBytes(details, p.Cert)
}
if p.InitiatorIndex != 0 { if p.InitiatorIndex != 0 {
details = protowire.AppendTag(details, fieldInitiatorIndex, protowire.VarintType) rest = protowire.AppendTag(rest, fieldInitiatorIndex, protowire.VarintType)
details = protowire.AppendVarint(details, uint64(p.InitiatorIndex)) rest = protowire.AppendVarint(rest, uint64(p.InitiatorIndex))
} }
if p.ResponderIndex != 0 { if p.ResponderIndex != 0 {
details = protowire.AppendTag(details, fieldResponderIndex, protowire.VarintType) rest = protowire.AppendTag(rest, fieldResponderIndex, protowire.VarintType)
details = protowire.AppendVarint(details, uint64(p.ResponderIndex)) rest = protowire.AppendVarint(rest, uint64(p.ResponderIndex))
} }
if p.Time != 0 { if p.Time != 0 {
details = protowire.AppendTag(details, fieldTime, protowire.VarintType) rest = protowire.AppendTag(rest, fieldTime, protowire.VarintType)
details = protowire.AppendVarint(details, p.Time) rest = protowire.AppendVarint(rest, p.Time)
} }
if p.CertVersion != 0 { if p.CertVersion != 0 {
details = protowire.AppendTag(details, fieldCertVersion, protowire.VarintType) rest = protowire.AppendTag(rest, fieldCertVersion, protowire.VarintType)
details = protowire.AppendVarint(details, uint64(p.CertVersion)) 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.AppendTag(out, 1, protowire.BytesType)
out = protowire.AppendBytes(out, details) out = protowire.AppendVarint(out, uint64(detailsLen))
if len(p.Cert) > 0 {
return out out = protowire.AppendTag(out, fieldCert, protowire.BytesType)
out = protowire.AppendBytes(out, p.Cert)
}
return append(out, rest...)
} }
// UnmarshalPayload decodes a protobuf-encoded NebulaHandshake message. // UnmarshalPayload decodes a protobuf-encoded NebulaHandshake message.