Skip to content

Commit b4b6a03

Browse files
committed
Remove unnecessary allocations in HopData encode/decode methods
By removing io.Writer and io.Reader abstraction from encode and decode methods, the following benchmars: old: func BenchmarkHopDataEncode(b *testing.B) { hd := new(HopData) for i := 0; i < b.N; i++ { hd.Encode(ioutil.Discard) } } type nopReader struct{} func (nopReader) Read(b []byte) (int, error) { return len(b), nil } func BenchmarkHopDataDecode(b *testing.B) { hd := new(HopData) src := nopReader{} for i := 0; i < b.N; i++ { hd.Decode(src) } } new: func BenchmarkHopDataEncode(b *testing.B) { hd := new(HopData) dst := make([]byte, hopDataSize) for i := 0; i < b.N; i++ { hd.Encode(dst) } } func BenchmarkHopDataDecode(b *testing.B) { hd := new(HopData) src := make([]byte, hopDataSize) for i := 0; i < b.N; i++ { err := hd.Decode(src) if err != nil { panic(err) } } } show the following insrease in performance: benchmark old ns/op new ns/op delta BenchmarkHopDataEncode-4 80.7 12.1 -85.01% benchmark old allocs new allocs delta BenchmarkHopDataEncode-4 3 0 -100.00% benchmark old bytes new bytes delta BenchmarkHopDataEncode-4 24 0 -100.00% benchmark old ns/op new ns/op delta BenchmarkHopDataDecode-4 218 9.77 -95.52% benchmark old allocs new allocs delta BenchmarkHopDataDecode-4 4 0 -100.00% benchmark old bytes new bytes delta BenchmarkHopDataDecode-4 56 0 -100.00%
1 parent c861375 commit b4b6a03

File tree

1 file changed

+26
-71
lines changed

1 file changed

+26
-71
lines changed

sphinx.go

+26-71
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@ import (
66
"crypto/hmac"
77
"crypto/sha256"
88
"encoding/binary"
9+
"errors"
910
"io"
10-
"io/ioutil"
1111
"math/big"
1212

1313
"github.com/aead/chacha20"
@@ -146,62 +146,33 @@ type HopData struct {
146146
HMAC [hmacSize]byte
147147
}
148148

149-
// Encode writes the serialized version of the target HopData into the passed
150-
// io.Writer.
151-
func (hd *HopData) Encode(w io.Writer) error {
152-
if _, err := w.Write([]byte{hd.Realm}); err != nil {
153-
return err
154-
}
155-
156-
if _, err := w.Write(hd.NextAddress[:]); err != nil {
157-
return err
158-
}
159-
160-
if err := binary.Write(w, binary.BigEndian, hd.ForwardAmount); err != nil {
161-
return err
162-
}
163-
164-
if err := binary.Write(w, binary.BigEndian, hd.OutgoingCltv); err != nil {
165-
return err
149+
// Encode writes the serialized version of the target HopData into the slice.
150+
func (hd *HopData) Encode(dst []byte) error {
151+
if len(dst) < hopDataSize {
152+
errors.New("destination is too small")
166153
}
167154

168-
if _, err := w.Write(paddingBytes[:]); err != nil {
169-
return err
170-
}
171-
172-
if _, err := w.Write(hd.HMAC[:]); err != nil {
173-
return err
174-
}
155+
dst[0] = hd.Realm
156+
copy(dst[1:], hd.NextAddress[:])
157+
binary.BigEndian.PutUint64(dst[1+addressSize:], hd.ForwardAmount)
158+
binary.BigEndian.PutUint32(dst[1+addressSize+8:], hd.OutgoingCltv)
159+
copy(dst[hopDataSize-hmacSize-padSize:], paddingBytes[:])
160+
copy(dst[hopDataSize-hmacSize:], hd.HMAC[:])
175161

176162
return nil
177163
}
178164

179-
// Decode deserializes the encoded HopData contained int he passed io.Reader
180-
// instance to the target empty HopData instance.
181-
func (hd *HopData) Decode(r io.Reader) error {
182-
if _, err := io.ReadFull(r, []byte{hd.Realm}); err != nil {
183-
return err
184-
}
185-
186-
if _, err := io.ReadFull(r, hd.NextAddress[:]); err != nil {
187-
return err
188-
}
189-
190-
if err := binary.Read(r, binary.BigEndian, &hd.ForwardAmount); err != nil {
191-
return err
192-
}
193-
194-
if err := binary.Read(r, binary.BigEndian, &hd.OutgoingCltv); err != nil {
195-
return err
196-
}
197-
198-
if _, err := io.CopyN(ioutil.Discard, r, padSize); err != nil {
199-
return err
165+
// Decode deserializes the encoded HopData contained in the passed slice.
166+
func (hd *HopData) Decode(src []byte) error {
167+
if len(src) < hopDataSize {
168+
errors.New("source is too small")
200169
}
201170

202-
if _, err := io.ReadFull(r, hd.HMAC[:]); err != nil {
203-
return err
204-
}
171+
hd.Realm = src[0]
172+
copy(hd.NextAddress[:], src[1:])
173+
hd.ForwardAmount = binary.BigEndian.Uint64(src[1+addressSize:])
174+
hd.OutgoingCltv = binary.BigEndian.Uint32(src[1+addressSize+8:])
175+
copy(hd.HMAC[:], src[hopDataSize-hmacSize:])
205176

206177
return nil
207178
}
@@ -292,9 +263,8 @@ func NewOnionPacket(paymentPath []*btcec.PublicKey, sessionKey *btcec.PrivateKey
292263
// Allocate zero'd out byte slices to store the final mix header packet
293264
// and the hmac for each hop.
294265
var (
295-
mixHeader [routingInfoSize]byte
296-
nextHmac [hmacSize]byte
297-
hopDataBuf bytes.Buffer
266+
mixHeader [routingInfoSize]byte
267+
nextHmac [hmacSize]byte
298268
)
299269

300270
// Now we compute the routing information for each hop, along with a
@@ -317,17 +287,16 @@ func NewOnionPacket(paymentPath []*btcec.PublicKey, sessionKey *btcec.PrivateKey
317287
streamBytes := generateCipherStream(rhoKey, numStreamBytes)
318288

319289
// Before we assemble the packet, we'll shift the current
320-
// mix-header to the write in order to make room for this next
290+
// mix-header to the right in order to make room for this next
321291
// per-hop data.
322-
rightShift(mixHeader[:], hopDataSize)
292+
copy(mixHeader[hopDataSize:], mixHeader[:routingInfoSize-hopDataSize])
323293

324294
// With the mix header right-shifted, we'll encode the current
325295
// hop data into a buffer we'll re-use during the packet
326296
// construction.
327-
if err := hopsData[i].Encode(&hopDataBuf); err != nil {
297+
if err := hopsData[i].Encode(mixHeader[:]); err != nil {
328298
return nil, err
329299
}
330-
copy(mixHeader[:], hopDataBuf.Bytes())
331300

332301
// Once the packet for this hop has been assembled, we'll
333302
// re-encrypt the packet by XOR'ing with a stream of bytes
@@ -346,8 +315,6 @@ func NewOnionPacket(paymentPath []*btcec.PublicKey, sessionKey *btcec.PrivateKey
346315
// prevent replay attacks.
347316
packet := append(mixHeader[:], assocData...)
348317
nextHmac = calcMac(muKey, packet)
349-
350-
hopDataBuf.Reset()
351318
}
352319

353320
return &OnionPacket{
@@ -358,18 +325,6 @@ func NewOnionPacket(paymentPath []*btcec.PublicKey, sessionKey *btcec.PrivateKey
358325
}, nil
359326
}
360327

361-
// Shift the byte-slice by the given number of bytes to the right and 0-fill
362-
// the resulting gap.
363-
func rightShift(slice []byte, num int) {
364-
for i := len(slice) - num - 1; i >= 0; i-- {
365-
slice[num+i] = slice[i]
366-
}
367-
368-
for i := 0; i < num; i++ {
369-
slice[i] = 0
370-
}
371-
}
372-
373328
// generateHeaderPadding derives the bytes for padding the mix header to ensure
374329
// it remains fixed sized throughout route transit. At each step, we add
375330
// 'hopSize' padding of zeroes, concatenate it to the previous filler, then
@@ -763,7 +718,7 @@ func processOnionPacket(onionPkt *OnionPacket,
763718
// out the per-hop data so we can derive the specified forwarding
764719
// instructions.
765720
var hopData HopData
766-
if err := hopData.Decode(bytes.NewReader(hopInfo[:])); err != nil {
721+
if err := hopData.Decode(hopInfo[:]); err != nil {
767722
return nil, err
768723
}
769724

0 commit comments

Comments
 (0)