From 78a53d3a882fc42beaf2b7c8c143f48a5f18ee28 Mon Sep 17 00:00:00 2001
From: Piotr Nojszewski <29924594+AeonSw4n@users.noreply.github.com>
Date: Thu, 7 Dec 2023 05:19:02 -0800
Subject: [PATCH] Add HandshakeController

---
 lib/pos_handshake_controller.go | 150 ++++++++++++++++++++++++++++++++
 1 file changed, 150 insertions(+)
 create mode 100644 lib/pos_handshake_controller.go

diff --git a/lib/pos_handshake_controller.go b/lib/pos_handshake_controller.go
new file mode 100644
index 000000000..8d06e2729
--- /dev/null
+++ b/lib/pos_handshake_controller.go
@@ -0,0 +1,150 @@
+package lib
+
+import (
+	"github.com/decred/dcrd/lru"
+	"github.com/deso-protocol/core/bls"
+	"github.com/golang/glog"
+)
+
+// TODO: Replace this
+func getActiveValidators() map[bls.PublicKey]*ValidatorEntry {
+	// TODO: replace with a getter to retrieve all active validators.
+	activeValidators := []*ValidatorEntry{}
+	allValidatorsMap := make(map[bls.PublicKey]*ValidatorEntry)
+	for _, validator := range activeValidators {
+		pk := validator.VotingPublicKey
+		if pk == nil {
+			continue
+		}
+		allValidatorsMap[*pk] = validator
+	}
+
+	return allValidatorsMap
+}
+
+type HandshakeController struct {
+	rniManager             *RemoteNodeIndexerManager
+	usedNonces             lru.Cache
+	protocolOnProofOfStake func() bool                              // TODO
+	getActiveValidators    func() map[bls.PublicKey]*ValidatorEntry // TODO
+}
+
+func NewHandshakeController(rniManager *RemoteNodeIndexerManager) *HandshakeController {
+
+	vm := &HandshakeController{
+		rniManager: rniManager,
+		usedNonces: lru.NewCache(1000),
+	}
+
+	return vm
+}
+
+func (hc *HandshakeController) handlePoSHandshakePeerMessage(origin *Peer, desoMsg DeSoMessage) {
+	if !hc.protocolOnProofOfStake() {
+		return
+	}
+
+	// Get the handshake information of this peer.
+	rn := hc.rniManager.GetRemoteNodeFromPeer(origin)
+	if rn == nil {
+		return
+	}
+
+	handshakeMetadata := rn.GetHandshakeMetadata()
+	// Make sure the peer is on the right proof of stake version
+	if handshakeMetadata.NegotiatedProtocolVersion() != ProtocolVersion2 {
+		// Disconnect the peer because we only accept validators running proof of stake.
+		hc.rniManager.Disconnect(rn)
+		//# cc.removeValidator(origin)
+		return
+	}
+
+	// Get all active validators and see if the peer is one of them.
+	activeValidators := hc.getActiveValidators()
+	validatorPk := handshakeMetadata.GetValidatorPublicKey()
+
+	// If there's already a validator connected with the same public key, disconnect the peer.
+	if _, ok := hc.rniManager.GetRemoteNodeIndexer().GetValidatorIndex().Get(validatorPk); ok {
+		hc.rniManager.Disconnect(rn)
+		return
+	}
+
+	// If the peer is not an active validator, there is nothing else to check so return.
+	if _, ok := activeValidators[validatorPk]; !ok {
+		return
+	}
+
+	// So we know this peer is an active validator. Add it to the validator index.
+	hc.rniManager.SetValidator(validatorPk, rn)
+}
+
+func (hc *HandshakeController) _handleHandshakePeerMessage(origin *Peer, desoMsg DeSoMessage) {
+	if desoMsg.GetMsgType() != MsgTypeHandshakePeer {
+		return
+	}
+
+	if hc.protocolOnProofOfStake() {
+		hc.handlePoSHandshakePeerMessage(origin, desoMsg)
+	}
+}
+
+func (hc *HandshakeController) _handleVersionMessage(origin *Peer, desoMsg DeSoMessage) {
+	if desoMsg.GetMsgType() != MsgTypeVersion {
+		return
+	}
+
+	rn := hc.rniManager.GetRemoteNodeFromPeer(origin)
+	if rn == nil {
+		return
+	}
+
+	var verMsg *MsgDeSoVersion
+	var ok bool
+	if verMsg, ok = desoMsg.(*MsgDeSoVersion); !ok {
+		hc.rniManager.Disconnect(rn)
+		return
+	}
+
+	// If we've sent this nonce before then return an error since this is
+	// a connection from ourselves.
+	msgNonce := verMsg.Nonce
+	if hc.usedNonces.Contains(msgNonce) {
+		hc.usedNonces.Delete(msgNonce)
+		glog.V(1).Infof("HandshakeController._handleVersionMessage: Requesting PeerDisconnect for id: (%v) "+
+			"nonce collision", origin.ID)
+		rn.Disconnect()
+		return
+	}
+
+	rn.HandleVersionMessage(verMsg, func(versionNonce uint64) {
+		hc.usedNonces.Add(versionNonce)
+	})
+}
+
+func (hc *HandshakeController) _handleVerackMessage(origin *Peer, desoMsg DeSoMessage) {
+	if desoMsg.GetMsgType() != MsgTypeVerack {
+		return
+	}
+
+	rn := hc.rniManager.GetRemoteNodeFromPeer(origin)
+	if rn == nil {
+		return
+	}
+
+	var vrkMsg *MsgDeSoVerack
+	var ok bool
+	if vrkMsg, ok = desoMsg.(*MsgDeSoVerack); !ok {
+		hc.rniManager.Disconnect(rn)
+		return
+	}
+
+	if !ok {
+		glog.V(1).Infof("HandshakeController._handleVerackMessage: Requesting PeerDisconnect for id: (%v) "+
+			"nonce not found for peer", origin.ID)
+		rn.Disconnect()
+		return
+	}
+
+	rn.HandleVerackMessage(vrkMsg)
+	return
+}