Extends cryptographic keys API, reduce frequency of OnMessageSent being called, improve cache locality of Kademlia routing table. (#274)

* client, node: have OnMessageSent called less frequently
kademlia/protocol, kademlia/table: switch from using a slice to a linked list for representing the routing table to make id updates more efficient, and add Last() to get the last entry of a bucket near a target public key
client, node, keys: add LoadKeysFromHex to load a private key from hex, add Signature type to represent cryptographic signatures, and add UnmarshalSignature which uses unsafe hackery to convert a byte slice into a signature
client: add missing check for length of overlay handshake

* client: use pkg-local constants instead of constants from ed25519

Author: iwasaki-kenta

client.go

@@ -9,7 +9,6 @@ import (
 	"encoding/hex"
 	"errors"
 	"fmt"
-	"github.com/oasislabs/ed25519"
 	"go.uber.org/zap"
 	"io"
 	"net"
@@ -411,7 +410,9 @@ func (c *Client) handshake() {
 
 	// Send our Ed25519 ephemeral public key and signature of the message '.__noise_handshake'.
 
-	if err := c.write(append(pub[:], sec.Sign([]byte(".__noise_handshake"))...)); err != nil {
+	signature := sec.Sign([]byte(".__noise_handshake"))
+
+	if err := c.write(append(pub[:], signature[:]...)); err != nil {
 		c.reportError(fmt.Errorf("failed to send session handshake: %w", err))
 		return
 	}
@@ -424,21 +425,21 @@ func (c *Client) handshake() {
 		return
 	}
 
-	if len(data) != ed25519.PublicKeySize+ed25519.SignatureSize {
+	if len(data) != SizePublicKey+SizeSignature {
 		c.reportError(fmt.Errorf("received invalid number of bytes opening a session: expected %d byte(s), but got %d byte(s)",
-			ed25519.PublicKeySize+ed25519.SignatureSize,
+			SizePublicKey+SizeSignature,
 			len(data),
 		))
 
 		return
 	}
 
 	var peerPublicKey PublicKey
-	copy(peerPublicKey[:], data[:ed25519.PublicKeySize])
+	copy(peerPublicKey[:], data[:SizePublicKey])
 
 	// Verify ownership of our peers Ed25519 public key by verifying the signature they sent.
 
-	if !peerPublicKey.Verify([]byte(".__noise_handshake"), data[ed25519.PublicKeySize:]) {
+	if !peerPublicKey.Verify([]byte(".__noise_handshake"), UnmarshalSignature(data[SizePublicKey:SizePublicKey+SizeSignature])) {
 		c.reportError(errors.New("could not verify session handshake"))
 		return
 	}
@@ -472,7 +473,8 @@ func (c *Client) handshake() {
 	// Send to our peer our overlay ID.
 
 	buf := c.node.id.Marshal()
-	buf = append(buf, c.node.Sign(append(buf, shared...))...)
+	signature = c.node.Sign(append(buf, shared...))
+	buf = append(buf, signature[:]...)
 
 	if err := c.write(buf); err != nil {
 		c.reportError(fmt.Errorf("failed to send session handshake: %w", err))
@@ -498,7 +500,16 @@ func (c *Client) handshake() {
 	buf = make([]byte, id.Size())
 	copy(buf, data)
 
-	if !id.ID.Verify(append(buf, shared...), data[len(buf):]) {
+	if len(data) != len(buf)+SizeSignature {
+		c.reportError(fmt.Errorf("received invalid number of bytes handshaking: expected %d byte(s), got %d byte(s)",
+			len(buf)+SizeSignature,
+			len(data),
+		))
+
+		return
+	}
+
+	if !id.ID.Verify(append(buf, shared...), UnmarshalSignature(data[len(buf):len(buf)+SizeSignature])) {
 		c.reportError(errors.New("overlay handshake signature is malformed"))
 		return
 	}
@@ -594,6 +605,14 @@ func (c *Client) writeLoop() {
 
 			break
 		}
+
+		for _, protocol := range c.node.protocols {
+			if protocol.OnMessageSent == nil {
+				continue
+			}
+
+			protocol.OnMessageSent(c)
+		}
 	}
 }
 

kademlia/protocol.go

@@ -103,8 +103,7 @@ func (p *Protocol) Ack(id noise.ID) {
 			return
 		}
 
-		bucket := p.table.Bucket(id.ID)
-		last := bucket[len(bucket)-1]
+		last := p.table.Last(id.ID)
 
 		ctx, cancel := context.WithTimeout(context.Background(), p.pingTimeout)
 		pong, err := p.node.RequestMessage(ctx, last.Address, Ping{})

kademlia/table.go

@@ -1,6 +1,7 @@
 package kademlia
 
 import (
+	"container/list"
 	"fmt"
 	"github.com/perlin-network/noise"
 	"sync"
@@ -10,7 +11,7 @@ import (
 type Table struct {
 	sync.RWMutex
 
-	entries [noise.SizePublicKey * 8][]noise.ID
+	entries [noise.SizePublicKey * 8]*list.List
 	self    noise.ID
 	size    int
 }
@@ -20,6 +21,10 @@ type Table struct {
 func NewTable(self noise.ID) *Table {
 	table := &Table{self: self}
 
+	for i := 0; i < len(table.entries); i++ {
+		table.entries[i] = list.New()
+	}
+
 	if _, err := table.Update(self); err != nil {
 		panic(err)
 	}
@@ -32,12 +37,27 @@ func (t *Table) Self() noise.ID {
 	return t.self
 }
 
-// Bucket returns all IDs in the bucket where target resides within.
+// Last returns the last id of the bucket where target resides within.
+func (t *Table) Last(target noise.PublicKey) noise.ID {
+	t.RLock()
+	defer t.RUnlock()
+
+	return t.entries[t.getBucketIndex(target)].Back().Value.(noise.ID)
+}
+
+// Bucket returns all entries of a bucket where target reside within.
 func (t *Table) Bucket(target noise.PublicKey) []noise.ID {
 	t.RLock()
 	defer t.RUnlock()
 
-	return t.entries[t.getBucketIndex(target)]
+	bucket := t.entries[t.getBucketIndex(target)]
+	entries := make([]noise.ID, 0, bucket.Len())
+
+	for e := bucket.Front(); e != nil; e = e.Next() {
+		entries = append(entries, e.Value.(noise.ID))
+	}
+
+	return entries
 }
 
 // Update attempts to insert the target node/peer ID into this routing table. If the bucket it was expected
@@ -52,17 +72,17 @@ func (t *Table) Update(target noise.ID) (bool, error) {
 	t.Lock()
 	defer t.Unlock()
 
-	idx := t.getBucketIndex(target.ID)
+	bucket := t.entries[t.getBucketIndex(target.ID)]
 
-	for i, id := range t.entries[idx] {
-		if id.ID == target.ID { // Found the target ID already inside the routing table.
-			t.entries[idx] = append(append([]noise.ID{target}, t.entries[idx][:i]...), t.entries[idx][i+1:]...)
+	for e := bucket.Front(); e != nil; e = e.Next() {
+		if e.Value.(noise.ID).ID == target.ID { // Found the target ID already inside the routing table.
+			bucket.MoveToFront(e)
 			return false, nil
 		}
 	}
 
-	if len(t.entries[idx]) < BucketSize { // The bucket is not yet under full capacity.
-		t.entries[idx] = append([]noise.ID{target}, t.entries[idx]...)
+	if bucket.Len() < BucketSize { // The bucket is not yet under full capacity.
+		bucket.PushFront(target)
 		t.size++
 		return true, nil
 	}
@@ -77,8 +97,10 @@ func (t *Table) Recorded(target noise.PublicKey) bool {
 	t.RLock()
 	defer t.RUnlock()
 
-	for _, id := range t.entries[t.getBucketIndex(target)] {
-		if id.ID == target {
+	bucket := t.entries[t.getBucketIndex(target)]
+
+	for e := bucket.Front(); e != nil; e = e.Next() {
+		if e.Value.(noise.ID).ID == target {
 			return true
 		}
 	}
@@ -92,11 +114,13 @@ func (t *Table) Delete(target noise.PublicKey) (noise.ID, bool) {
 	t.Lock()
 	defer t.Unlock()
 
-	idx := t.getBucketIndex(target)
+	bucket := t.entries[t.getBucketIndex(target)]
+
+	for e := bucket.Front(); e != nil; e = e.Next() {
+		id := e.Value.(noise.ID)
 
-	for i, id := range t.entries[idx] {
 		if id.ID == target {
-			t.entries[idx] = append(t.entries[idx][:i], t.entries[idx][i+1:]...)
+			bucket.Remove(e)
 			t.size--
 			return id, true
 		}
@@ -111,10 +135,12 @@ func (t *Table) DeleteByAddress(target string) (noise.ID, bool) {
 	t.Lock()
 	defer t.Unlock()
 
-	for i, bucket := range t.entries {
-		for j, id := range bucket {
+	for _, bucket := range t.entries {
+		for e := bucket.Front(); e != nil; e = e.Next() {
+			id := e.Value.(noise.ID)
+
 			if id.Address == target {
-				t.entries[i] = append(t.entries[i][:j], t.entries[i][j+1:]...)
+				bucket.Remove(e)
 				t.size--
 				return id, true
 			}
@@ -133,8 +159,10 @@ func (t *Table) Peers() []noise.ID {
 func (t *Table) FindClosest(target noise.PublicKey, k int) []noise.ID {
 	var closest []noise.ID
 
-	f := func(bucket []noise.ID) {
-		for _, id := range bucket {
+	f := func(bucket *list.List) {
+		for e := bucket.Front(); e != nil; e = e.Next() {
+			id := e.Value.(noise.ID)
+
 			if id.ID != target {
 				closest = append(closest, id)
 			}
@@ -175,8 +203,8 @@ func (t *Table) Entries() []noise.ID {
 	entries := make([]noise.ID, 0, t.size)
 
 	for _, bucket := range t.entries {
-		for _, id := range bucket {
-			entries = append(entries, id)
+		for e := bucket.Front(); e != nil; e = e.Next() {
+			entries = append(entries, e.Value.(noise.ID))
 		}
 	}
 

keys.go

@@ -3,8 +3,11 @@ package noise
 import (
 	"encoding/hex"
 	"encoding/json"
+	"fmt"
 	"github.com/oasislabs/ed25519"
 	"io"
+	"reflect"
+	"unsafe"
 )
 
 const (
@@ -13,6 +16,9 @@ const (
 
 	// SizePrivateKey is the size in bytes of a nodes/peers private key.
 	SizePrivateKey = ed25519.PrivateKeySize
+
+	// SizeSignature is the size in bytes of a cryptographic signature.
+	SizeSignature = ed25519.SignatureSize
 )
 
 type (
@@ -21,6 +27,9 @@ type (
 
 	// PrivateKey is the default node/peer private key type.
 	PrivateKey [SizePrivateKey]byte
+
+	// Signature is the default node/peer cryptographic signature type.
+	Signature [SizeSignature]byte
 )
 
 var (
@@ -29,10 +38,13 @@ var (
 
 	// ZeroPrivateKey is the zero-value for a node/peer private key.
 	ZeroPrivateKey PrivateKey
+
+	// ZeroSignature is the zero-value for a cryptographic signature.
+	ZeroSignature Signature
 )
 
 // GenerateKeys randomly generates a new pair of cryptographic keys. Nil may be passed to rand in order to use
-// crypto/rand by default. It throws an error if rand is invalid.
+// crypto/rand by default. It returns an error if rand is invalid.
 func GenerateKeys(rand io.Reader) (publicKey PublicKey, privateKey PrivateKey, err error) {
 	pub, priv, err := ed25519.GenerateKey(rand)
 	if err != nil {
@@ -45,9 +57,30 @@ func GenerateKeys(rand io.Reader) (publicKey PublicKey, privateKey PrivateKey, e
 	return publicKey, privateKey, nil
 }
 
+// LoadKeysFromHex loads a private key from a hex string. It returns an error if secretHex is not hex-encoded or is
+// an invalid number of bytes. In the case of the latter error, the error is wrapped as io.ErrUnexpectedEOF. Calling
+// this function performs 1 allocation.
+func LoadKeysFromHex(secretHex string) (PrivateKey, error) {
+	secret, err := hex.DecodeString(secretHex)
+	if err != nil {
+		return ZeroPrivateKey, fmt.Errorf("private key provided in hex failed to be decoded: %w", err)
+	}
+
+	if len(secret) != SizePrivateKey {
+		return ZeroPrivateKey, fmt.Errorf("got private key of %d byte(s), but expected %d byte(s): %w",
+			len(secret), SizePrivateKey, io.ErrUnexpectedEOF,
+		)
+	}
+
+	var privateKey PrivateKey
+	copy(privateKey[:], secret)
+
+	return privateKey, nil
+}
+
 // Verify returns true if the cryptographic signature of data is representative of this public key.
-func (k PublicKey) Verify(data, signature []byte) bool {
-	return ed25519.Verify(k[:], data, signature)
+func (k PublicKey) Verify(data []byte, signature Signature) bool {
+	return ed25519.Verify(k[:], data, signature[:])
 }
 
 // String returns the hexadecimal representation of this public key.
@@ -61,8 +94,8 @@ func (k PublicKey) MarshalJSON() ([]byte, error) {
 }
 
 // Sign uses this private key to sign data and return its cryptographic signature as a slice of bytes.
-func (k PrivateKey) Sign(data []byte) []byte {
-	return ed25519.Sign(k[:], data)
+func (k PrivateKey) Sign(data []byte) Signature {
+	return UnmarshalSignature(ed25519.Sign(k[:], data))
 }
 
 // String returns the hexadecimal representation of this private key.
@@ -72,7 +105,7 @@ func (k PrivateKey) String() string {
 
 // MarshalJSON returns the hexadecimal representation of this private key in JSON. It should never throw an error.
 func (k PrivateKey) MarshalJSON() ([]byte, error) {
-	return json.Marshal(hex.EncodeToString(k[:]))
+	return json.Marshal(k.String())
 }
 
 // Public returns the public key associated to this private key.
@@ -82,3 +115,20 @@ func (k PrivateKey) Public() PublicKey {
 
 	return publicKey
 }
+
+// String returns the hexadecimal representation of this signature.
+func (s Signature) String() string {
+	return hex.EncodeToString(s[:])
+}
+
+// MarshalJSON returns the hexadecimal representation of this signature in JSON. It should never throw an error.
+func (s Signature) MarshalJSON() ([]byte, error) {
+	return json.Marshal(s.String())
+}
+
+// UnmarshalSignature decodes data into a Signature instance. It panics if data is not of expected length by instilling
+// a bound check hint to the compiler. It uses unsafe hackery to zero-alloc convert data into a Signature.
+func UnmarshalSignature(data []byte) Signature {
+	_ = data[SizeSignature-1]
+	return *(*Signature)(unsafe.Pointer((*reflect.SliceHeader)(unsafe.Pointer(&data)).Data))
+}

mod.go

@@ -39,7 +39,7 @@ type Protocol struct {
 	// OnPingFailed is called whenever any attempt by a node to dial a peer at addr fails.
 	OnPingFailed func(addr string, err error)
 
-	// OnMessageSent is called whenever a message or request is successfully sent to a peer.
+	// OnMessageSent is called whenever bytes of a message or request or response have been flushed/sent to a peer.
 	OnMessageSent func(client *Client)
 
 	// OnMessageRecv is called whenever a message or response is received from a peer.