keep track of recent nonces with a heap and a map instead of a fixed-size bitmask

src/yggdrasil/session.go

@@ -6,6 +6,7 @@ package yggdrasil
 
 import (
 	"bytes"
+	"container/heap"
 	"errors"
 	"sync"
 	"time"
@@ -15,6 +16,25 @@ import (
 	"github.com/yggdrasil-network/yggdrasil-go/src/util"
 )
 
+// Duration that we keep track of old nonces per session, to allow some out-of-order packet delivery
+const nonceWindow = time.Second
+
+// A heap of nonces, used with a map[nonce]time to allow out-of-order packets a little time to arrive without rejecting them
+// Less is backwards so the oldest node is the highest priority for Pop
+type nonceHeap []crypto.BoxNonce
+
+func (h nonceHeap) Len() int            { return len(h) }
+func (h nonceHeap) Less(i, j int) bool  { return h[i].Minus(&h[j]) > 0 }
+func (h nonceHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
+func (h *nonceHeap) Push(x interface{}) { *h = append(*h, x.(crypto.BoxNonce)) }
+func (h *nonceHeap) Pop() interface{} {
+	l := len(*h)
+	var n crypto.BoxNonce
+	n, *h = (*h)[l-1], (*h)[:l-1]
+	return n
+}
+func (h nonceHeap) peek() *crypto.BoxNonce { return &h[len(h)-1] }
+
 // All the information we know about an active session.
 // This includes coords, permanent and ephemeral keys, handles and nonces, various sorts of timing information for timeout and maintenance, and some metadata for the admin API.
 type sessionInfo struct {
@@ -31,7 +51,8 @@ type sessionInfo struct {
 	theirHandle    crypto.Handle                 //
 	myHandle       crypto.Handle                 //
 	theirNonce     crypto.BoxNonce               //
-	theirNonceMask uint64                  //
+	theirNonceHeap nonceHeap                     // priority queue to keep track of the lowest nonce we recently accepted
+	theirNonceMap  map[crypto.BoxNonce]time.Time // time we added each nonce to the heap
 	myNonce        crypto.BoxNonce               //
 	theirMTU       uint16                        //
 	myMTU          uint16                        //
@@ -87,7 +108,8 @@ func (s *sessionInfo) update(p *sessionPing) bool {
 		s.theirHandle = p.Handle
 		s.sharedSesKey = *crypto.GetSharedKey(&s.mySesPriv, &s.theirSesPub)
 		s.theirNonce = crypto.BoxNonce{}
-		s.theirNonceMask = 0
+		s.theirNonceHeap = nil
+		s.theirNonceMap = make(map[crypto.BoxNonce]time.Time)
 	}
 	if p.MTU >= 1280 || p.MTU == 0 {
 		s.theirMTU = p.MTU
@@ -400,27 +422,40 @@ func (sinfo *sessionInfo) getMTU() uint16 {
 // Checks if a packet's nonce is recent enough to fall within the window of allowed packets, and not already received.
 func (sinfo *sessionInfo) nonceIsOK(theirNonce *crypto.BoxNonce) bool {
 	// The bitmask is to allow for some non-duplicate out-of-order packets
-	diff := theirNonce.Minus(&sinfo.theirNonce)
-	if diff > 0 {
+	if theirNonce.Minus(&sinfo.theirNonce) > 0 {
+		// This is newer than the newest nonce we've seen
 		return true
 	}
-	return ^sinfo.theirNonceMask&(0x01<<uint64(-diff)) != 0
+	if len(sinfo.theirNonceHeap) > 0 {
+		if theirNonce.Minus(sinfo.theirNonceHeap.peek()) > 0 {
+			if _, isIn := sinfo.theirNonceMap[*theirNonce]; !isIn {
+				// This nonce is recent enough that we keep track of older nonces, but it's not one we've seen yet
+				return true
+			}
+		}
+	}
+	return false
 }
 
 // Updates the nonce mask by (possibly) shifting the bitmask and setting the bit corresponding to this nonce to 1, and then updating the most recent nonce
 func (sinfo *sessionInfo) updateNonce(theirNonce *crypto.BoxNonce) {
-	// Shift nonce mask if needed
-	// Set bit
-	diff := theirNonce.Minus(&sinfo.theirNonce)
-	if diff > 0 {
-		// This nonce is newer, so shift the window before setting the bit, and update theirNonce in the session info.
-		sinfo.theirNonceMask <<= uint64(diff)
-		sinfo.theirNonceMask &= 0x01
-		sinfo.theirNonce = *theirNonce
-	} else {
-		// This nonce is older, so set the bit but do not shift the window.
-		sinfo.theirNonceMask &= 0x01 << uint64(-diff)
+	// Start with some cleanup
+	for len(sinfo.theirNonceHeap) > 64 {
+		if time.Since(sinfo.theirNonceMap[*sinfo.theirNonceHeap.peek()]) < nonceWindow {
+			// This nonce is still fairly new, so keep it around
+			break
 		}
+		// TODO? reallocate the map in some cases, to free unused map space?
+		delete(sinfo.theirNonceMap, *sinfo.theirNonceHeap.peek())
+		heap.Pop(&sinfo.theirNonceHeap)
+	}
+	if theirNonce.Minus(&sinfo.theirNonce) > 0 {
+		// This nonce is the newest we've seen, so make a note of that
+		sinfo.theirNonce = *theirNonce
+	}
+	// Add it to the heap/map so we know not to allow it again
+	heap.Push(&sinfo.theirNonceHeap, *theirNonce)
+	sinfo.theirNonceMap[*theirNonce] = time.Now()
 }
 
 // Resets all sessions to an uninitialized state.