device: optimize message encoding

Optimize message encoding by eliminating binary.Write (which internally
uses reflection) in favour of hand-rolled encoding.

This is companion to 9e7529c3d2d0c54f4d5384c01645a9279e4740ae.

Synthetic benchmark:

    var packetSink []byte
    func BenchmarkMessageInitiationMarshal(b *testing.B) {
        var msg MessageInitiation
        b.Run("binary.Write", func(b *testing.B) {
            b.ReportAllocs()
            for range b.N {
                var buf [MessageInitiationSize]byte
                writer := bytes.NewBuffer(buf[:0])
                _ = binary.Write(writer, binary.LittleEndian, msg)
                packetSink = writer.Bytes()
            }
        })
        b.Run("binary.Encode", func(b *testing.B) {
            b.ReportAllocs()
            for range b.N {
                packet := make([]byte, MessageInitiationSize)
                _, _ = binary.Encode(packet, binary.LittleEndian, msg)
                packetSink = packet
            }
        })
        b.Run("marshal", func(b *testing.B) {
            b.ReportAllocs()
            for range b.N {
                packet := make([]byte, MessageInitiationSize)
                _ = msg.marshal(packet)
                packetSink = packet
            }
        })
    }

Results:
                                             │      -      │
                                             │   sec/op    │
    MessageInitiationMarshal/binary.Write-8    1.337µ ± 0%
    MessageInitiationMarshal/binary.Encode-8   1.242µ ± 0%
    MessageInitiationMarshal/marshal-8         53.05n ± 1%

                                             │     -      │
                                             │    B/op    │
    MessageInitiationMarshal/binary.Write-8    368.0 ± 0%
    MessageInitiationMarshal/binary.Encode-8   160.0 ± 0%
    MessageInitiationMarshal/marshal-8         160.0 ± 0%

                                             │     -      │
                                             │ allocs/op  │
    MessageInitiationMarshal/binary.Write-8    3.000 ± 0%
    MessageInitiationMarshal/binary.Encode-8   1.000 ± 0%
    MessageInitiationMarshal/marshal-8         1.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 5f713ee5ed5581fb6be1e10fd365104b8cc9b73c..5cf1702b68a1748b653e7ab8e6b0cc19648c37d8 100644 (file)
--- a/device/noise-protocol.go
+++ b/device/noise-protocol.go
@@ -134,6 +134,22 @@ func (msg *MessageInitiation) unmarshal(b []byte) error {
        return nil
 }
 
+func (msg *MessageInitiation) marshal(b []byte) error {
+       if len(b) != MessageInitiationSize {
+               return errMessageLengthMismatch
+       }
+
+       binary.LittleEndian.PutUint32(b, msg.Type)
+       binary.LittleEndian.PutUint32(b[4:], msg.Sender)
+       copy(b[8:], msg.Ephemeral[:])
+       copy(b[8+len(msg.Ephemeral):], msg.Static[:])
+       copy(b[8+len(msg.Ephemeral)+len(msg.Static):], msg.Timestamp[:])
+       copy(b[8+len(msg.Ephemeral)+len(msg.Static)+len(msg.Timestamp):], msg.MAC1[:])
+       copy(b[8+len(msg.Ephemeral)+len(msg.Static)+len(msg.Timestamp)+len(msg.MAC1):], msg.MAC2[:])
+
+       return nil
+}
+
 func (msg *MessageResponse) unmarshal(b []byte) error {
        if len(b) != MessageResponseSize {
                return errMessageLengthMismatch
@@ -150,6 +166,22 @@ func (msg *MessageResponse) unmarshal(b []byte) error {
        return nil
 }
 
+func (msg *MessageResponse) marshal(b []byte) error {
+       if len(b) != MessageResponseSize {
+               return errMessageLengthMismatch
+       }
+
+       binary.LittleEndian.PutUint32(b, msg.Type)
+       binary.LittleEndian.PutUint32(b[4:], msg.Sender)
+       binary.LittleEndian.PutUint32(b[8:], msg.Receiver)
+       copy(b[12:], msg.Ephemeral[:])
+       copy(b[12+len(msg.Ephemeral):], msg.Empty[:])
+       copy(b[12+len(msg.Ephemeral)+len(msg.Empty):], msg.MAC1[:])
+       copy(b[12+len(msg.Ephemeral)+len(msg.Empty)+len(msg.MAC1):], msg.MAC2[:])
+
+       return nil
+}
+
 func (msg *MessageCookieReply) unmarshal(b []byte) error {
        if len(b) != MessageCookieReplySize {
                return errMessageLengthMismatch
@@ -163,6 +195,19 @@ func (msg *MessageCookieReply) unmarshal(b []byte) error {
        return nil
 }
 
+func (msg *MessageCookieReply) marshal(b []byte) error {
+       if len(b) != MessageCookieReplySize {
+               return errMessageLengthMismatch
+       }
+
+       binary.LittleEndian.PutUint32(b, msg.Type)
+       binary.LittleEndian.PutUint32(b[4:], msg.Receiver)
+       copy(b[8:], msg.Nonce[:])
+       copy(b[8+len(msg.Nonce):], msg.Cookie[:])
+
+       return nil
+}
+
 type Handshake struct {
        state                     handshakeState
        mutex                     sync.RWMutex

diff --git a/device/send.go b/device/send.go
index 38f55c2a4805b0608af61d1ff90554fb8d3d7b36..ff8f7da507434464fbb52989a41656377ec93b2c 100644 (file)
--- a/device/send.go
+++ b/device/send.go
@@ -6,7 +6,6 @@
 package device
 
 import (
-       "bytes"
        "encoding/binary"
        "errors"
        "net"
@@ -124,10 +123,8 @@ func (peer *Peer) SendHandshakeInitiation(isRetry bool) error {
                return err
        }
 
-       var buf [MessageInitiationSize]byte
-       writer := bytes.NewBuffer(buf[:0])
-       binary.Write(writer, binary.LittleEndian, msg)
-       packet := writer.Bytes()
+       packet := make([]byte, MessageInitiationSize)
+       _ = msg.marshal(packet)
        peer.cookieGenerator.AddMacs(packet)
 
        peer.timersAnyAuthenticatedPacketTraversal()
@@ -155,10 +152,8 @@ func (peer *Peer) SendHandshakeResponse() error {
                return err
        }
 
-       var buf [MessageResponseSize]byte
-       writer := bytes.NewBuffer(buf[:0])
-       binary.Write(writer, binary.LittleEndian, response)
-       packet := writer.Bytes()
+       packet := make([]byte, MessageResponseSize)
+       _ = response.marshal(packet)
        peer.cookieGenerator.AddMacs(packet)
 
        err = peer.BeginSymmetricSession()
@@ -189,11 +184,11 @@ func (device *Device) SendHandshakeCookie(initiatingElem *QueueHandshakeElement)
                return err
        }
 
-       var buf [MessageCookieReplySize]byte
-       writer := bytes.NewBuffer(buf[:0])
-       binary.Write(writer, binary.LittleEndian, reply)
+       packet := make([]byte, MessageCookieReplySize)
+       _ = reply.marshal(packet)
        // TODO: allocation could be avoided
-       device.net.bind.Send([][]byte{writer.Bytes()}, initiatingElem.endpoint)
+       device.net.bind.Send([][]byte{packet}, initiatingElem.endpoint)
+
        return nil
 }