diff --git a/runelite-client/src/main/java/net/runelite/client/account/SessionManager.java b/runelite-client/src/main/java/net/runelite/client/account/SessionManager.java index 510fd99e27..54de7b774e 100644 --- a/runelite-client/src/main/java/net/runelite/client/account/SessionManager.java +++ b/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 ScheduledExecutorService executor; + private final ConfigManager configManager; + 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)) - { - if (wsclient != null) - { - wsclient.close(); - } - - wsclient = new WSClient(eventBus, executor, session); - wsclient.connect(); - } + wsClient.changeSession(session.getUuid()); } accountSession = session; @@ -165,11 +155,7 @@ public class SessionManager private void closeSession() { - if (wsclient != null) - { - wsclient.close(); - wsclient = null; - } + wsClient.close(); if (accountSession == null) { diff --git a/runelite-client/src/main/java/net/runelite/client/account/WSClient.java b/runelite-client/src/main/java/net/runelite/client/ws/WSClient.java similarity index 68% rename from runelite-client/src/main/java/net/runelite/client/account/WSClient.java rename to runelite-client/src/main/java/net/runelite/client/ws/WSClient.java index 16ba8b36a6..279794862f 100644 --- a/runelite-client/src/main/java/net/runelite/client/account/WSClient.java +++ b/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> 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 message) + { + if (messages.add(message)) + { + gson = WebsocketGsonFactory.build(WebsocketGsonFactory.factory(messages)); + } + } + + public void unregisterMessage(final Class 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); - log.debug("Got message: {}", message); + final WebsocketMessage message = gson.fromJson(text, WebsocketMessage.class); + if (message.isParty() && !(message instanceof PartyMessage)) + { + // spoofed message? + return; + } + log.debug("Got message: {}", message); eventBus.post(message); }