device: reduce RoutineHandshake allocations

Reduce allocations by eliminating byte reader, hand-rolled decoding and
reusing message structs.

Synthetic benchmark:

    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)
            }
        })
    }

Results:
                                         │      -      │
                                         │   sec/op    │
MessageInitiationUnmarshal/binary.Read-8   1.508µ ± 2%
MessageInitiationUnmarshal/unmarshal-8     12.66n ± 2%

                                         │      -       │
                                         │     B/op     │
MessageInitiationUnmarshal/binary.Read-8   208.0 ± 0%
MessageInitiationUnmarshal/unmarshal-8     0.000 ± 0%

                                         │      -       │
                                         │  allocs/op   │
MessageInitiationUnmarshal/binary.Read-8   2.000 ± 0%
MessageInitiationUnmarshal/unmarshal-8     0.000 ± 0%

Signed-off-by: Alexander Yastrebov <yastrebov.alex@gmail.com>
Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

diff --git a/device/noise-protocol.go b/device/noise-protocol.go
index b72acf85d5e426ffb26b5c4b9d2de7678125fd50..12368ec62e68b0c4cea0a09ae5cedb35212da49c 100644 (file)
--- 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 c7b6c87fc594f91c5f888dcb73eb711e1b926b1b..13929577ee8f6844a897ec45f3c02d97f821f925 100644 (file)
--- 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