From 041ea4ac1a0edbdd0417b3aee516fb7a915fbc32 Mon Sep 17 00:00:00 2001
From: Adam Gibson <ekaggata@gmail.com>
Date: Thu, 7 Apr 2022 17:07:27 +0100
Subject: [PATCH] Choose directory node based on where nicks seen.

Original testing setup of onion message channels just sent messages on
the first directory node peer in our list. This fixes that TODO.
After this commit, we keep a map of which directory nodes we have seen a
given Joinmarket nick on, and we keep track of whether those directory
nodes are available (connected) or not. Then when we want to privmsg
that nick, if we do not currently have a direct connection, we choose
randomly from the set of directory nodes on which that nick has been
seen, and which are currently live/connected.
---
 jmdaemon/jmdaemon/onionmc.py | 49 +++++++++++++++++++++++++++++++++---
 1 file changed, 45 insertions(+), 4 deletions(-)

diff --git a/jmdaemon/jmdaemon/onionmc.py b/jmdaemon/jmdaemon/onionmc.py
index ad0059f..72dd210 100644
--- a/jmdaemon/jmdaemon/onionmc.py
+++ b/jmdaemon/jmdaemon/onionmc.py
@@ -3,6 +3,7 @@ from jmdaemon.protocol import COMMAND_PREFIX, JM_VERSION
 from jmbase import get_log,  JM_APP_NAME, JMHiddenService, stop_reactor
 import json
 import copy
+import random
 from typing import Callable, Union, Tuple, List
 from twisted.internet import reactor, task, protocol
 from twisted.protocols import basic
@@ -113,6 +114,9 @@ class OnionCustomMessageDecodingError(Exception):
 class InvalidLocationStringError(Exception):
     pass
 
+class OnionDirectoryPeerNotFound(Exception):
+    pass
+
 class OnionCustomMessage(object):
     """ Encapsulates the messages passed over the wire
     to and from other onion peers
@@ -583,6 +587,14 @@ class OnionDirectoryPeer(OnionPeer):
         except OnionPeerConnectionError:
             reactor.callLater(self.delay, self.try_to_connect)
 
+    def register_connection(self) -> None:
+        self.messagechannel.update_directory_map(self, connected=True)
+        super().register_connection()
+
+    def register_disconnection(self) -> None:
+        self.messagechannel.update_directory_map(self, connected=False)
+        super().register_disconnection()
+
 
 class OnionMessageChannel(MessageChannel):
     """ Sends messages to other nodes of the same type over Tor
@@ -683,6 +695,14 @@ class OnionMessageChannel(MessageChannel):
         # the rpc connection calls are not using twisted)
         self.wait_for_directories_loop = None
 
+        # this dict plays the same role as `active_channels` in `MessageChannelCollection`.
+        # it has structure {nick: set(),..} where set() has elements that are dicts:
+        # {OnionPeer: bool}.
+        # Entries get updated with changing connection status of directories,
+        # allowing us to decide where to send each message we want to send when we have no
+        # direct connection.
+        self.active_directories = {}
+
     def info_callback(self, msg: str) -> None:
         log.info(msg)
 
@@ -783,11 +803,10 @@ class OnionMessageChannel(MessageChannel):
             log.debug("Privmsg peer: {} but don't have peerid; "
                      "sending via directory.".format(nick))
             try:
-                # TODO change this to redundant or switching
-                peer_sendable = self.get_connected_directory_peers()[0]
-            except Exception as e:
+                peer_sendable = self.get_directory_for_nick(nick)
+            except OnionDirectoryPeerNotFound:
                 log.warn("Failed to send privmsg because no "
-                "directory peer is connected. Error: {}".format(repr(e)))
+                "directory peer is connected.")
                 return
         self._send(peer_sendable, encoded_privmsg)
 
@@ -961,6 +980,28 @@ class OnionMessageChannel(MessageChannel):
         except Exception as e:
             log.debug("Invalid Joinmarket message: {}, error was: {}".format(
                 msgval, repr(e)))
+        # add the nick to the directories map, whether pubmsg or privmsg, but
+        # only if it passed the above syntax Exception catch:
+        if peer.directory and not self.self_as_peer.directory:
+            if from_nick not in self.active_directories:
+                self.active_directories[from_nick] = {}
+            self.active_directories[from_nick][peer] = True
+
+    def update_directory_map(self, p: OnionDirectoryPeer, connected: bool) -> None:
+        nicks = []
+        for nick in self.active_directories:
+            if p in self.active_directories[nick]:
+                nicks.append(nick)
+        for nick in nicks:
+            self.active_directories[nick][p] = connected
+
+    def get_directory_for_nick(self, nick: str) -> OnionDirectoryPeer:
+        if nick not in self.active_directories:
+            raise OnionDirectoryPeerNotFound
+        adn = self.active_directories[nick]
+        if len(adn) == 0:
+            raise OnionDirectoryPeerNotFound
+        return random.choice([x for x in list(adn) if adn[x] is True])
 
     def forward_pubmsg_to_peers(self, msg: str, from_nick: str) -> None:
         """ Used by directory nodes currently. Takes a received