From 688f215dfd5c8f81afdfdc133d90f1654095b7cd Mon Sep 17 00:00:00 2001
From: catbref <misc-github@talk2dom.com>
Date: Mon, 21 Dec 2020 12:25:58 +0000
Subject: [PATCH] Allow PresenceType filtering on presence websocket via
 presenceType query param

---
 .../api/websocket/PresenceWebSocket.java      | 33 +++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/main/java/org/qortal/api/websocket/PresenceWebSocket.java b/src/main/java/org/qortal/api/websocket/PresenceWebSocket.java
index 27bc603f..5fd191f0 100644
--- a/src/main/java/org/qortal/api/websocket/PresenceWebSocket.java
+++ b/src/main/java/org/qortal/api/websocket/PresenceWebSocket.java
@@ -79,6 +79,9 @@ public class PresenceWebSocket extends ApiWebSocket implements Listener {
 	/** Outer map key is PresenceType (enum), inner map key is public key in base58, inner map value is timestamp */
 	private static final Map<PresenceType, Map<String, Long>> currentEntries = Collections.synchronizedMap(new EnumMap<>(PresenceType.class));
 
+	/** (Optional) PresenceType used for filtering by that Session. */
+	private static final Map<Session, PresenceType> sessionPresenceTypes = Collections.synchronizedMap(new HashMap<>());
+
 	@Override
 	public void configure(WebSocketServletFactory factory) {
 		factory.register(PresenceWebSocket.class);
@@ -118,17 +121,40 @@ public class PresenceWebSocket extends ApiWebSocket implements Listener {
 		List<PresenceInfo> presenceInfo = Collections.singletonList(new PresenceInfo(presenceType, pubKey58, computedTimestamp));
 
 		// Notify sessions
-		for (Session session : getSessions())
-			sendPresenceInfo(session, presenceInfo);
+		for (Session session : getSessions()) {
+			PresenceType sessionPresenceType = sessionPresenceTypes.get(session);
+
+			if (sessionPresenceType == null || sessionPresenceType == presenceType)
+				sendPresenceInfo(session, presenceInfo);
+		}
 	}
 
 	@OnWebSocketConnect
 	@Override
 	public void onWebSocketConnect(Session session) {
+		Map<String, List<String>> queryParams = session.getUpgradeRequest().getParameterMap();
+		List<String> presenceTypes = queryParams.get("presenceType");
+
+		// We only support ONE presenceType
+		String presenceTypeName = presenceTypes == null || presenceTypes.isEmpty() ? null : presenceTypes.get(0);
+
+		PresenceType presenceType = presenceTypeName == null ? null : PresenceType.fromString(presenceTypeName);
+
+		// Make sure that if caller does give a presenceType, that it is a valid/known one.
+		if (presenceTypeName != null && presenceType == null) {
+			session.close(4003, "unknown presenceType: " + presenceTypeName);
+			return;
+		}
+
+		// Save session's requested PresenceType, if given
+		if (presenceType != null)
+			sessionPresenceTypes.put(session, presenceType);
+
 		List<PresenceInfo> presenceInfo;
 
 		synchronized (currentEntries) {
 			presenceInfo = currentEntries.entrySet().stream()
+					.filter(entry -> presenceType == null ? true : entry.getKey() == presenceType)
 					.flatMap(entry -> entry.getValue().entrySet().stream().map(innerEntry -> new PresenceInfo(entry.getKey(), innerEntry.getKey(), innerEntry.getValue())))
 					.collect(Collectors.toList());
 		}
@@ -144,6 +170,9 @@ public class PresenceWebSocket extends ApiWebSocket implements Listener {
 	@OnWebSocketClose
 	@Override
 	public void onWebSocketClose(Session session, int statusCode, String reason) {
+		// clean up
+		sessionPresenceTypes.remove(session);
+
 		super.onWebSocketClose(session, statusCode, reason);
 	}