From af55375359409ed825a77352d019720e02d50fcb Mon Sep 17 00:00:00 2001 From: Alexander Yastrebov Date: Thu, 26 Dec 2024 20:36:53 +0100 Subject: [PATCH 1/2] device: reduce RoutineHandshake allocations Reduce allocations by eliminating byte reader, hand-rolled decoding and reusing message structs. Signed-off-by: Alexander Yastrebov --- device/noise-protocol.go | 48 ++++++++++++++++++++++++++++++++++++++++ device/receive.go | 10 +++------ 2 files changed, 51 insertions(+), 7 deletions(-) diff --git a/device/noise-protocol.go b/device/noise-protocol.go index b72acf85d..12368ec62 100644 --- a/device/noise-protocol.go +++ b/device/noise-protocol.go @@ -6,6 +6,7 @@ package device import ( + "encoding/binary" "errors" "fmt" "sync" @@ -115,6 +116,53 @@ type MessageCookieReply struct { Cookie [blake2s.Size128 + poly1305.TagSize]byte } +var errMessageTooShort = errors.New("message too short") + +func (msg *MessageInitiation) unmarshal(b []byte) error { + if len(b) < MessageInitiationSize { + return errMessageTooShort + } + + msg.Type = binary.LittleEndian.Uint32(b) + msg.Sender = binary.LittleEndian.Uint32(b[4:]) + copy(msg.Ephemeral[:], b[8:]) + copy(msg.Static[:], b[8+len(msg.Ephemeral):]) + copy(msg.Timestamp[:], b[8+len(msg.Ephemeral)+len(msg.Static):]) + copy(msg.MAC1[:], b[8+len(msg.Ephemeral)+len(msg.Static)+len(msg.Timestamp):]) + copy(msg.MAC2[:], b[8+len(msg.Ephemeral)+len(msg.Static)+len(msg.Timestamp)+len(msg.MAC1):]) + + return nil +} + +func (msg *MessageResponse) unmarshal(b []byte) error { + if len(b) < MessageResponseSize { + return errMessageTooShort + } + + msg.Type = binary.LittleEndian.Uint32(b) + msg.Sender = binary.LittleEndian.Uint32(b[4:]) + msg.Receiver = binary.LittleEndian.Uint32(b[8:]) + copy(msg.Ephemeral[:], b[12:]) + copy(msg.Empty[:], b[12+len(msg.Ephemeral):]) + copy(msg.MAC1[:], b[12+len(msg.Ephemeral)+len(msg.Empty):]) + copy(msg.MAC2[:], b[12+len(msg.Ephemeral)+len(msg.Empty)+len(msg.MAC1):]) + + return nil +} + +func (msg *MessageCookieReply) unmarshal(b []byte) error { + if len(b) < MessageCookieReplySize { + return errMessageTooShort + } + + msg.Type = binary.LittleEndian.Uint32(b) + msg.Receiver = binary.LittleEndian.Uint32(b[4:]) + copy(msg.Nonce[:], b[8:]) + copy(msg.Cookie[:], b[8+len(msg.Nonce):]) + + return nil +} + type Handshake struct { state handshakeState mutex sync.RWMutex diff --git a/device/receive.go b/device/receive.go index c7b6c87fc..13929577e 100644 --- a/device/receive.go +++ b/device/receive.go @@ -6,7 +6,6 @@ package device import ( - "bytes" "encoding/binary" "errors" "net" @@ -287,8 +286,7 @@ func (device *Device) RoutineHandshake(id int) { // unmarshal packet var reply MessageCookieReply - reader := bytes.NewReader(elem.packet) - err := binary.Read(reader, binary.LittleEndian, &reply) + err := reply.unmarshal(elem.packet) if err != nil { device.log.Verbosef("Failed to decode cookie reply") goto skip @@ -353,8 +351,7 @@ func (device *Device) RoutineHandshake(id int) { // unmarshal var msg MessageInitiation - reader := bytes.NewReader(elem.packet) - err := binary.Read(reader, binary.LittleEndian, &msg) + err := msg.unmarshal(elem.packet) if err != nil { device.log.Errorf("Failed to decode initiation message") goto skip @@ -386,8 +383,7 @@ func (device *Device) RoutineHandshake(id int) { // unmarshal var msg MessageResponse - reader := bytes.NewReader(elem.packet) - err := binary.Read(reader, binary.LittleEndian, &msg) + err := msg.unmarshal(elem.packet) if err != nil { device.log.Errorf("Failed to decode response message") goto skip From 39abcab4054f6fadae52388e33e4d100ba5896e1 Mon Sep 17 00:00:00 2001 From: Alexander Yastrebov Date: Fri, 9 May 2025 14:53:48 +0200 Subject: [PATCH 2/2] device: add BenchmarkMessageInitiationUnmarshal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit goos: linux goarch: amd64 pkg: golang.zx2c4.com/wireguard/device │ - │ │ sec/op │ MessageInitiationUnmarshal/binary.Read-8 1.508µ ± 2% MessageInitiationUnmarshal/unmarshal-8 12.66n ± 2% geomean 138.1n │ - │ │ B/op │ MessageInitiationUnmarshal/binary.Read-8 208.0 ± 0% MessageInitiationUnmarshal/unmarshal-8 0.000 ± 0% geomean ¹ ¹ summaries must be >0 to compute geomean │ - │ │ allocs/op │ MessageInitiationUnmarshal/binary.Read-8 2.000 ± 0% MessageInitiationUnmarshal/unmarshal-8 0.000 ± 0% geomean ¹ ¹ summaries must be >0 to compute geomean Signed-off-by: Alexander Yastrebov --- device/noise-protocol_test.go | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 device/noise-protocol_test.go diff --git a/device/noise-protocol_test.go b/device/noise-protocol_test.go new file mode 100644 index 000000000..7812e325d --- /dev/null +++ b/device/noise-protocol_test.go @@ -0,0 +1,33 @@ +package device + +import ( + "bytes" + "encoding/binary" + "testing" +) + +var msgSink MessageInitiation + +func BenchmarkMessageInitiationUnmarshal(b *testing.B) { + packet := make([]byte, MessageInitiationSize) + reader := bytes.NewReader(packet) + err := binary.Read(reader, binary.LittleEndian, &msgSink) + if err != nil { + b.Fatal(err) + } + + b.Run("binary.Read", func(b *testing.B) { + b.ReportAllocs() + for range b.N { + reader := bytes.NewReader(packet) + _ = binary.Read(reader, binary.LittleEndian, &msgSink) + } + }) + + b.Run("unmarshal", func(b *testing.B) { + b.ReportAllocs() + for range b.N { + _ = msgSink.unmarshal(packet) + } + }) +}