Split WSClient from SessionManager

Co-authored-by: Tomas Slusny <slusnucky@gmail.com>

runelite-client/src/main/java/net/runelite/client/account/SessionManager.java

@@ -31,7 +31,6 @@ import java.io.FileWriter;
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.time.Instant;
-import java.util.concurrent.ScheduledExecutorService;
 import javax.inject.Inject;
 import javax.inject.Singleton;
 import lombok.Getter;
@@ -43,6 +42,7 @@ import net.runelite.client.config.ConfigManager;
 import net.runelite.client.eventbus.EventBus;
 import net.runelite.client.eventbus.Subscribe;
 import net.runelite.client.util.LinkBrowser;
+import net.runelite.client.ws.WSClient;
 import net.runelite.http.api.account.AccountClient;
 import net.runelite.http.api.account.OAuthResponse;
 import net.runelite.http.api.ws.messages.LoginResponse;
@@ -52,22 +52,21 @@ import net.runelite.http.api.ws.messages.LoginResponse;
 public class SessionManager
 {
 	private static final File SESSION_FILE = new File(RuneLite.RUNELITE_DIR, "session");
-	private WSClient wsclient;
 
 	@Getter
 	private AccountSession accountSession;
 
 	private final EventBus eventBus;
-	private ConfigManager configManager;
+	private final ConfigManager configManager;
-	private ScheduledExecutorService executor;
+	private final WSClient wsClient;
 	private final AccountClient loginClient = new AccountClient();
 
 	@Inject
-	public SessionManager(ConfigManager configManager, EventBus eventBus, ScheduledExecutorService executor)
+	private SessionManager(ConfigManager configManager, EventBus eventBus, WSClient wsClient)
 	{
 		this.configManager = configManager;
 		this.eventBus = eventBus;
-		this.executor = executor;
+		this.wsClient = wsClient;
 		eventBus.register(this);
 	}
 
@@ -136,19 +135,10 @@ public class SessionManager
 	 */
 	private void openSession(AccountSession session, boolean openSocket)
 	{
-		// If the ws session already exists, don't need to do anything
+		// Change session on the websocket
 		if (openSocket)
 		{
-			if (wsclient == null || !wsclient.checkSession(session))
+			wsClient.changeSession(session.getUuid());
-			{
-				if (wsclient != null)
-				{
-					wsclient.close();
-				}
-
-				wsclient = new WSClient(eventBus, executor, session);
-				wsclient.connect();
-			}
 		}
 
 		accountSession = session;
@@ -165,11 +155,7 @@ public class SessionManager
 
 	private void closeSession()
 	{
-		if (wsclient != null)
+		wsClient.close();
-		{
-			wsclient.close();
-			wsclient = null;
-		}
 
 		if (accountSession == null)
 		{

runelite-client/src/main/java/net/runelite/client/ws/WSClient.java

@@ -22,14 +22,20 @@
  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */
-package net.runelite.client.account;
+package net.runelite.client.ws;
 
 import com.google.gson.Gson;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Objects;
+import java.util.UUID;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
+import javax.inject.Inject;
+import javax.inject.Singleton;
 import lombok.extern.slf4j.Slf4j;
 import net.runelite.client.eventbus.EventBus;
 import net.runelite.http.api.RuneLiteAPI;
@@ -37,6 +43,7 @@ import net.runelite.http.api.ws.WebsocketGsonFactory;
 import net.runelite.http.api.ws.WebsocketMessage;
 import net.runelite.http.api.ws.messages.Handshake;
 import net.runelite.http.api.ws.messages.Ping;
+import net.runelite.http.api.ws.messages.party.PartyMessage;
 import okhttp3.OkHttpClient;
 import okhttp3.Request;
 import okhttp3.Response;
@@ -44,31 +51,61 @@ import okhttp3.WebSocket;
 import okhttp3.WebSocketListener;
 
 @Slf4j
+@Singleton
 public class WSClient extends WebSocketListener implements AutoCloseable
 {
 	private static final Duration PING_TIME = Duration.ofSeconds(30);
-	private static final Gson GSON = WebsocketGsonFactory.build();
 
 	private final OkHttpClient client = new OkHttpClient();
 	private final EventBus eventBus;
 	private final ScheduledFuture pingFuture;
-	private final AccountSession session;
+	private final Collection<Class<? extends WebsocketMessage>> messages = new HashSet<>();
+
+	private volatile Gson gson;
+	private UUID sessionId;
 	private WebSocket webSocket;
 
-	WSClient(EventBus eventBus, ScheduledExecutorService executor, AccountSession session)
+	@Inject
+	private WSClient(EventBus eventBus, ScheduledExecutorService executor)
 	{
 		this.eventBus = eventBus;
-		this.session = session;
 		this.pingFuture = executor.scheduleWithFixedDelay(this::ping, PING_TIME.getSeconds(), PING_TIME.getSeconds(), TimeUnit.SECONDS);
+		this.gson = WebsocketGsonFactory.build(WebsocketGsonFactory.factory(messages));
 	}
 
-	boolean checkSession(AccountSession session)
+	public boolean sessionExists()
 	{
-		return session.equals(this.session);
+		return sessionId != null;
 	}
 
-	void connect()
+	public void changeSession(UUID sessionId)
 	{
+		if (Objects.equals(sessionId, this.sessionId))
+		{
+			return;
+		}
+
+		if (webSocket != null)
+		{
+			close();
+			webSocket = null;
+		}
+
+		this.sessionId = sessionId;
+
+		if (sessionId != null)
+		{
+			connect();
+		}
+	}
+
+	private void connect()
+	{
+		if (sessionId == null)
+		{
+			throw new IllegalStateException("Cannot connect with no session id");
+		}
+
 		Request request = new Request.Builder()
 			.url(RuneLiteAPI.getWsEndpoint())
 			.build();
@@ -76,7 +113,7 @@ public class WSClient extends WebSocketListener implements AutoCloseable
 		webSocket = client.newWebSocket(request, this);
 
 		Handshake handshake = new Handshake();
-		handshake.setSession(session.getUuid());
+		handshake.setSession(sessionId);
 		send(handshake);
 	}
 
@@ -93,29 +130,38 @@ public class WSClient extends WebSocketListener implements AutoCloseable
 		send(ping);
 	}
 
-	private void send(WebsocketMessage message)
+	public void registerMessage(final Class<? extends WebsocketMessage> message)
+	{
+		if (messages.add(message))
+		{
+			gson = WebsocketGsonFactory.build(WebsocketGsonFactory.factory(messages));
+		}
+	}
+
+	public void unregisterMessage(final Class<? extends WebsocketMessage> message)
+	{
+		if (messages.remove(message))
+		{
+			gson = WebsocketGsonFactory.build(WebsocketGsonFactory.factory(messages));
+		}
+	}
+
+	public void send(WebsocketMessage message)
 	{
 		if (webSocket == null)
 		{
 			log.debug("Reconnecting to server");
-
 			connect();
 		}
 
-		String json = GSON.toJson(message, WebsocketMessage.class);
+		final String json = gson.toJson(message, WebsocketMessage.class);
 		webSocket.send(json);
-
 		log.debug("Sent: {}", json);
 	}
 
 	@Override
 	public void close()
 	{
-		if (pingFuture != null)
-		{
-			pingFuture.cancel(true);
-		}
-
 		if (webSocket != null)
 		{
 			webSocket.close(1000, null);
@@ -131,9 +177,14 @@ public class WSClient extends WebSocketListener implements AutoCloseable
 	@Override
 	public void onMessage(WebSocket webSocket, String text)
 	{
-		WebsocketMessage message = GSON.fromJson(text, WebsocketMessage.class);
+		final WebsocketMessage message = gson.fromJson(text, WebsocketMessage.class);
-		log.debug("Got message: {}", message);
+		if (message.isParty() && !(message instanceof PartyMessage))
+		{
+			// spoofed message?
+			return;
+		}
 
+		log.debug("Got message: {}", message);
 		eventBus.post(message);
 	}