diff --git a/LaunchServer/source/response/ServerSocketHandler.java b/LaunchServer/source/response/ServerSocketHandler.java index a4c380d..ee801e4 100644 --- a/LaunchServer/source/response/ServerSocketHandler.java +++ b/LaunchServer/source/response/ServerSocketHandler.java @@ -12,6 +12,7 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiConsumer; import java.util.function.Consumer; +import java.util.function.Predicate; import launcher.LauncherAPI; import launcher.helper.CommonHelper; @@ -34,6 +35,7 @@ // API private final Map customResponses = new ConcurrentHashMap<>(2); private volatile BiConsumer connectListener; + private volatile Predicate preConnectListener; private volatile Consumer disconnectListener; public ServerSocketHandler(LaunchServer server) { @@ -70,7 +72,12 @@ // Listen for incoming connections while (serverSocket.isBound()) { - threadPool.execute(new ResponseThread(server, serverSocket.accept())); + Socket socket = serverSocket.accept(); + if (preConnectListener != null && !preConnectListener.test(socket)) { + IOHelper.close(socket); + continue; + } + threadPool.execute(new ResponseThread(server, socket)); } } catch (IOException e) { // Ignore error after close/rebind @@ -104,6 +111,11 @@ this.disconnectListener = disconnectListener; } + @LauncherAPI + public void setPreConnectListener(Predicate preConnectListener) { + this.preConnectListener = preConnectListener; + } + /*package*/ void onConnected(Socket socket, Request.Type type) { if (connectListener != null) { connectListener.accept(socket, type);