void/amneziawg-go
2
0
Fork 0
mirror of https://github.com/amnezia-vpn/amneziawg-go.git synced 2025-07-30 08:32:50 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

device: optimize message encoding

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

This is companion to 9e7529c3d2.

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>
This commit is contained in:
Alexander Yastrebov 2025-05-17 11:34:30 +02:00 • committed by Jason A. Donenfeld
parent 256bcbd70d
commit 264889f0bb
2 changed files with 53 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
device/noise-protocol.go
View file

@ -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

21
device/send.go
View file

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