From 934cd1d5116528a786ab4b66194f290402283186 Mon Sep 17 00:00:00 2001
From: catbref <misc-github@talk2dom.com>
Date: Thu, 10 Dec 2020 17:08:32 +0000
Subject: [PATCH] Add support for preferred blockchain to
 /websockets/crosschain/tradebot via foreignBlockchain query param

---
 .../api/websocket/TradeBotWebSocket.java      | 43 ++++++++++++++++---
 1 file changed, 36 insertions(+), 7 deletions(-)

diff --git a/src/main/java/org/qortal/api/websocket/TradeBotWebSocket.java b/src/main/java/org/qortal/api/websocket/TradeBotWebSocket.java
index 4a8bd348..55969c6b 100644
--- a/src/main/java/org/qortal/api/websocket/TradeBotWebSocket.java
+++ b/src/main/java/org/qortal/api/websocket/TradeBotWebSocket.java
@@ -16,6 +16,7 @@ import org.eclipse.jetty.websocket.api.annotations.OnWebSocketMessage;
 import org.eclipse.jetty.websocket.api.annotations.WebSocket;
 import org.eclipse.jetty.websocket.servlet.WebSocketServletFactory;
 import org.qortal.controller.tradebot.TradeBot;
+import org.qortal.crosschain.SupportedBlockchain;
 import org.qortal.data.crosschain.TradeBotData;
 import org.qortal.event.Event;
 import org.qortal.event.EventBus;
@@ -32,6 +33,8 @@ public class TradeBotWebSocket extends ApiWebSocket implements Listener {
 	/** Cache of trade-bot entry states, keyed by trade-bot entry's "trade private key" (base58) */
 	private static final Map<String, Integer> PREVIOUS_STATES = new HashMap<>();
 
+	private static final Map<Session, String> sessionBlockchain = Collections.synchronizedMap(new HashMap<>());
+
 	@Override
 	public void configure(WebSocketServletFactory factory) {
 		factory.register(TradeBotWebSocket.class);
@@ -68,27 +71,50 @@ public class TradeBotWebSocket extends ApiWebSocket implements Listener {
 		}
 
 		List<TradeBotData> tradeBotEntries = Collections.singletonList(tradeBotData);
-		for (Session session : getSessions())
-			sendEntries(session, tradeBotEntries);
+
+		for (Session session : getSessions()) {
+			// Only send if this session has this/no preferred blockchain
+			String preferredBlockchain = sessionBlockchain.get(session);
+
+			if (preferredBlockchain == null || preferredBlockchain.equals(tradeBotData.getForeignBlockchain()))
+				sendEntries(session, tradeBotEntries);
+		}
 	}
 
 	@OnWebSocketConnect
 	@Override
 	public void onWebSocketConnect(Session session) {
+		Map<String, List<String>> queryParams = session.getUpgradeRequest().getParameterMap();
+
+		List<String> foreignBlockchains = queryParams.get("foreignBlockchain");
+		final String foreignBlockchain = foreignBlockchains == null ? null : foreignBlockchains.get(0);
+
+		// Make sure blockchain (if any) is valid
+		if (foreignBlockchain != null && SupportedBlockchain.fromString(foreignBlockchain) == null) {
+			session.close(4003, "unknown blockchain: " + foreignBlockchain);
+			return;
+		}
+
+		// save session's preferred blockchain (if any)
+		sessionBlockchain.put(session, foreignBlockchain);
+
 		// Send all known trade-bot entries
 		try (final Repository repository = RepositoryManager.getRepository()) {
 			List<TradeBotData> tradeBotEntries = repository.getCrossChainRepository().getAllTradeBotData();
-			if (tradeBotEntries == null) {
-				session.close(4001, "repository issue fetching trade-bot entries");
-				return;
-			}
+
+			// Optional filtering
+			if (foreignBlockchain != null)
+				tradeBotEntries = tradeBotEntries.stream()
+						.filter(tradeBotData -> tradeBotData.getForeignBlockchain().equals(foreignBlockchain))
+						.collect(Collectors.toList());
 
 			if (!sendEntries(session, tradeBotEntries)) {
 				session.close(4002, "websocket issue");
 				return;
 			}
 		} catch (DataException e) {
-			// No output this time
+			session.close(4001, "repository issue fetching trade-bot entries");
+			return;
 		}
 
 		super.onWebSocketConnect(session);
@@ -97,6 +123,9 @@ public class TradeBotWebSocket extends ApiWebSocket implements Listener {
 	@OnWebSocketClose
 	@Override
 	public void onWebSocketClose(Session session, int statusCode, String reason) {
+		// clean up
+		sessionBlockchain.remove(session);
+
 		super.onWebSocketClose(session, statusCode, reason);
 	}