Protocol sequences: Add support for bidirectional responses.

protocol/Sequence.go

@@ -6,11 +6,16 @@ Author:     Peter Kleissner
 This code caches and verifies message sequences. Sequence numbers are valid on a peer level, independent of which network connection was used.
 They can be used to map incoming response messages to previous outgoing requests. The remote peer ID is used together with a consecutive sequence number as unique key.
 
+Sequences are created for:
+* Unidirectional messages (responses are one-way)
+* Bidirectional messages (responses may be sent back and forth). They use a different key namespace since remote sequence numbers could collide with locally created ones.
+
 Advantages:
 * This secures against replay and poisoning attacks.
 * If used correctly it can also deduplicate messages (which occurs when 2 peers have multiple registered connections to each other but none are active and subsequent fallback to broadcast).
 * The round-trip time can be measured and used to determine the connection quality.
 * (future) It can be used to detect missed and lost replies.
+
 */
 
 package protocol
@@ -43,6 +48,10 @@ type SequenceExpiry struct {
 	expires        time.Time   // When the sequence expires. This can be extended on the fly!
 	counter        int         // How many replies used the sequence. Multiple Response messages may be returned for a single Announcement one.
 	Data           interface{} // Optional high-level data associated with the sequence
+	// bidirectional sequences only
+	bidirectional  bool          // Whether this sequence is used in a bidirectional way
+	timeout        time.Duration // Timeout for receiving the next message
+	invalidateFunc func()        // The invalidation callback is in case a sequence collision or expiration invalidates the sequence.
 }
 
 // NewSequenceManager creates a new sequence manager. The ReplyTimeout is in seconds. The expiration function is started immediately.
@@ -67,12 +76,26 @@ func (manager *SequenceManager) autoDeleteExpired() {
 		for key, sequence := range manager.sequences {
 			if sequence.expires.Before(now) {
 				delete(manager.sequences, key)
+
+				if sequence.invalidateFunc != nil {
+					go sequence.invalidateFunc()
+				}
 			}
 		}
 		manager.Unlock()
 	}
 }
 
+// sequence2Key creates the lookup key of a sequence for a peer.
+// Since bidirectional sequence numbers may be created from either side (remote or local peer), it does not share a namespace with unidirectional sequence numbers.
+func sequence2Key(bidirectional bool, publicKey *btcec.PublicKey, sequenceNumber uint32) (key string) {
+	if !bidirectional {
+		return "u" + string(publicKey.SerializeCompressed()) + strconv.FormatUint(uint64(sequenceNumber), 10)
+	} else {
+		return "b" + string(publicKey.SerializeCompressed()) + strconv.FormatUint(uint64(sequenceNumber), 10)
+	}
+}
+
 // NewSequence returns a new sequence and registers it. messageSequence must point to the variable holding the continuous next sequence number.
 // Use only for Announcement and Ping messages.
 func (manager *SequenceManager) NewSequence(publicKey *btcec.PublicKey, messageSequence *uint32, data interface{}) (info *SequenceExpiry) {
@@ -84,7 +107,7 @@ func (manager *SequenceManager) NewSequence(publicKey *btcec.PublicKey, messageS
 	}
 
 	// Add the sequence to the list. Sequences are unique enough that collisions are unlikely and negligible.
-	key := string(publicKey.SerializeCompressed()) + strconv.FormatUint(uint64(info.SequenceNumber), 10)
+	key := sequence2Key(false, publicKey, info.SequenceNumber)
 	manager.Lock()
 	manager.sequences[key] = info
 	manager.Unlock()
@@ -102,7 +125,7 @@ func (manager *SequenceManager) ArbitrarySequence(publicKey *btcec.PublicKey, da
 	}
 
 	// Add the sequence to the list. Sequences are unique enough that collisions are unlikely and negligible.
-	key := string(publicKey.SerializeCompressed()) + strconv.FormatUint(uint64(info.SequenceNumber), 10)
+	key := sequence2Key(false, publicKey, info.SequenceNumber)
 	manager.Lock()
 	manager.sequences[key] = info
 	manager.Unlock()
@@ -112,7 +135,7 @@ func (manager *SequenceManager) ArbitrarySequence(publicKey *btcec.PublicKey, da
 
 // ValidateSequence validates the sequence number of an incoming message. It will set raw.sequence if valid.
 func (manager *SequenceManager) ValidateSequence(publicKey *btcec.PublicKey, sequenceNumber uint32, invalidate, extendValidity bool) (sequenceInfo *SequenceExpiry, valid bool, rtt time.Duration) {
-	key := string(publicKey.SerializeCompressed()) + strconv.FormatUint(uint64(sequenceNumber), 10)
+	key := sequence2Key(false, publicKey, sequenceNumber)
 
 	manager.Lock()
 	defer manager.Unlock()
@@ -144,9 +167,96 @@ func (manager *SequenceManager) ValidateSequence(publicKey *btcec.PublicKey, seq
 
 // InvalidateSequence invalidates the sequence number.
 func (manager *SequenceManager) InvalidateSequence(publicKey *btcec.PublicKey, sequenceNumber uint32) {
-	key := string(publicKey.SerializeCompressed()) + strconv.FormatUint(uint64(sequenceNumber), 10)
+	key := sequence2Key(false, publicKey, sequenceNumber)
 
 	manager.Lock()
 	delete(manager.sequences, key)
 	manager.Unlock()
 }
+
+// ---- bidirectional sequences ----
+
+// RegisterSequenceBi registers a bidirectional sequence initiated by a remote peer. The caller must specify the timeout (which will be reset every time a new message appears in this sequence).
+// This is needed for bidirectional responses to accept subsequent incoming messages from the remote peer.
+func (manager *SequenceManager) RegisterSequenceBi(publicKey *btcec.PublicKey, sequenceNumber uint32, data interface{}, timeout time.Duration, invalidateFunc func()) (info *SequenceExpiry) {
+	info = &SequenceExpiry{
+		SequenceNumber: sequenceNumber,
+		created:        time.Now(),
+		expires:        time.Now().Add(timeout),
+		timeout:        timeout,
+		invalidateFunc: invalidateFunc,
+		Data:           data,
+	}
+
+	// Before registering the sequence, check if there is a collision. If yes, invalidate the original one.
+	key := sequence2Key(true, publicKey, info.SequenceNumber)
+	manager.Lock()
+	existingSequence := manager.sequences[key]
+	manager.sequences[key] = info
+	manager.Unlock()
+
+	// Call the invalidate function if there is a collision.
+	if existingSequence != nil && existingSequence.invalidateFunc != nil {
+		go existingSequence.invalidateFunc()
+	}
+
+	return
+}
+
+// NewSequenceBi returns a new bidirectional sequence and registers it. messageSequence must point to the variable holding the continuous next sequence number.
+func (manager *SequenceManager) NewSequenceBi(publicKey *btcec.PublicKey, messageSequence *uint32, data interface{}, timeout time.Duration, invalidateFunc func()) (info *SequenceExpiry) {
+	info = &SequenceExpiry{
+		created:        time.Now(),
+		expires:        time.Now().Add(timeout),
+		bidirectional:  true,
+		timeout:        timeout,
+		invalidateFunc: invalidateFunc,
+		Data:           data,
+	}
+
+	manager.Lock()
+	defer manager.Unlock()
+
+	// The likelihood of a collision is low but not impossible.
+	for n := 0; n < 10000; n++ {
+		info.SequenceNumber = atomic.AddUint32(messageSequence, 1)
+
+		key := sequence2Key(true, publicKey, info.SequenceNumber)
+		if infoE := manager.sequences[key]; infoE == nil {
+			manager.sequences[key] = info
+			return info
+		}
+	}
+
+	return nil
+}
+
+// ValidateSequenceBi validates the sequence number of an incoming message. It will set raw.sequence if valid.
+func (manager *SequenceManager) ValidateSequenceBi(publicKey *btcec.PublicKey, sequenceNumber uint32, isLast bool) (sequenceInfo *SequenceExpiry, valid bool, rtt time.Duration) {
+	key := sequence2Key(true, publicKey, sequenceNumber)
+
+	manager.Lock()
+	defer manager.Unlock()
+
+	// lookup the sequence
+	sequence, ok := manager.sequences[key]
+	if !ok {
+		return nil, false, rtt
+	}
+
+	// Initial reply: Store latest roundtrip time.
+	if sequence.counter == 0 {
+		rtt = time.Since(sequence.created)
+	}
+
+	sequence.counter++
+
+	// invalidate the sequence immediately?
+	if isLast {
+		delete(manager.sequences, key)
+	} else {
+		sequence.expires = time.Now().Add(sequence.timeout)
+	}
+
+	return sequence, sequence.expires.After(time.Now()), rtt
+}