From 184911f76e9d5164d6402f46e656e61971cbc4e9 Mon Sep 17 00:00:00 2001
From: stasatdaglabs <39559713+stasatdaglabs@users.noreply.github.com>
Date: Tue, 17 Nov 2020 18:05:14 +0200
Subject: [PATCH] [NOD-1547] Make the SendAddresses flow not one-time for the
 sake of DNSSeeder (#1104)

* [NOD-1547] Make the SendAddresses flow not one-time for the sake of DNSSeeder.

* [NOD-1547] Add all special commands to chooseRouteForCommand.
---
 .../flows/addressexchange/sendaddresses.go    | 29 +++++++++++--------
 app/protocol/protocol.go                      |  2 +-
 .../network/netadapter/standalone/routes.go   | 16 +++++++++-
 3 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/app/protocol/flows/addressexchange/sendaddresses.go b/app/protocol/flows/addressexchange/sendaddresses.go
index 340640459..f7a9cdd17 100644
--- a/app/protocol/flows/addressexchange/sendaddresses.go
+++ b/app/protocol/flows/addressexchange/sendaddresses.go
@@ -16,20 +16,25 @@ type SendAddressesContext interface {
 
 // SendAddresses sends addresses to a peer that requests it.
 func SendAddresses(context SendAddressesContext, incomingRoute *router.Route, outgoingRoute *router.Route) error {
-	message, err := incomingRoute.Dequeue()
-	if err != nil {
-		return err
-	}
+	for {
+		message, err := incomingRoute.Dequeue()
+		if err != nil {
+			return err
+		}
 
-	_, ok := message.(*appmessage.MsgRequestAddresses)
-	if !ok {
-		return protocolerrors.Errorf(true, "unexpected message. "+
-			"Expected: %s, got: %s", appmessage.CmdRequestAddresses, message.Command())
-	}
-	addresses := context.AddressManager().Addresses()
-	msgAddresses := appmessage.NewMsgAddresses(shuffleAddresses(addresses))
+		_, ok := message.(*appmessage.MsgRequestAddresses)
+		if !ok {
+			return protocolerrors.Errorf(true, "unexpected message. "+
+				"Expected: %s, got: %s", appmessage.CmdRequestAddresses, message.Command())
+		}
+		addresses := context.AddressManager().Addresses()
+		msgAddresses := appmessage.NewMsgAddresses(shuffleAddresses(addresses))
 
-	return outgoingRoute.Enqueue(msgAddresses)
+		err = outgoingRoute.Enqueue(msgAddresses)
+		if err != nil {
+			return err
+		}
+	}
 }
 
 // shuffleAddresses randomizes the given addresses sent if there are more than the maximum allowed in one message.
diff --git a/app/protocol/protocol.go b/app/protocol/protocol.go
index 523018796..28c0ee772 100644
--- a/app/protocol/protocol.go
+++ b/app/protocol/protocol.go
@@ -118,7 +118,7 @@ func (m *Manager) registerAddressFlows(router *routerpkg.Router, isStopping *uin
 	outgoingRoute := router.OutgoingRoute()
 
 	return []*flow{
-		m.registerOneTimeFlow("SendAddresses", router, []appmessage.MessageCommand{appmessage.CmdRequestAddresses}, isStopping, errChan,
+		m.registerFlow("SendAddresses", router, []appmessage.MessageCommand{appmessage.CmdRequestAddresses}, isStopping, errChan,
 			func(incomingRoute *routerpkg.Route, peer *peerpkg.Peer) error {
 				return addressexchange.SendAddresses(m.context, incomingRoute, outgoingRoute)
 			},
diff --git a/infrastructure/network/netadapter/standalone/routes.go b/infrastructure/network/netadapter/standalone/routes.go
index d2f7d0083..0f1dc830a 100644
--- a/infrastructure/network/netadapter/standalone/routes.go
+++ b/infrastructure/network/netadapter/standalone/routes.go
@@ -25,7 +25,8 @@ type Routes struct {
 func (r *Routes) WaitForMessageOfType(command appmessage.MessageCommand, timeout time.Duration) (appmessage.Message, error) {
 	timeoutTime := time.Now().Add(timeout)
 	for {
-		message, err := r.IncomingRoute.DequeueWithTimeout(timeoutTime.Sub(time.Now()))
+		route := r.chooseRouteForCommand(command)
+		message, err := route.DequeueWithTimeout(timeoutTime.Sub(time.Now()))
 		if err != nil {
 			return nil, errors.Wrapf(err, "error waiting for message of type %s", command)
 		}
@@ -35,6 +36,19 @@ func (r *Routes) WaitForMessageOfType(command appmessage.MessageCommand, timeout
 	}
 }
 
+func (r *Routes) chooseRouteForCommand(command appmessage.MessageCommand) *router.Route {
+	switch command {
+	case appmessage.CmdVersion, appmessage.CmdVerAck:
+		return r.handshakeRoute
+	case appmessage.CmdRequestAddresses, appmessage.CmdAddresses:
+		return r.addressesRoute
+	case appmessage.CmdPing:
+		return r.pingRoute
+	default:
+		return r.IncomingRoute
+	}
+}
+
 // WaitForDisconnect waits for a disconnect up to `timeout`, skipping all messages received while waiting
 func (r *Routes) WaitForDisconnect(timeout time.Duration) error {
 	timeoutTime := time.Now().Add(timeout)