diff --git a/plugins/MinikuraVelocity/src/main/kotlin/cafe/kirameki/minikuraVelocity/Main.kt b/plugins/MinikuraVelocity/src/main/kotlin/cafe/kirameki/minikuraVelocity/Main.kt index 8183b07..5964fce 100644 --- a/plugins/MinikuraVelocity/src/main/kotlin/cafe/kirameki/minikuraVelocity/Main.kt +++ b/plugins/MinikuraVelocity/src/main/kotlin/cafe/kirameki/minikuraVelocity/Main.kt @@ -111,7 +111,7 @@ class Main @Inject constructor(private val logger: Logger, private val server: P commandManager.register(refreshCommandMeta, refreshCommand) commandManager.register(migrateCommandMeta, migrateCommand) - val connectionHandler = ProxyTransferHandler(servers, logger) + val connectionHandler = ProxyTransferHandler(servers, logger, acceptingTransfers) server.eventManager.register(this, connectionHandler) Executors.newSingleThreadExecutor().submit { diff --git a/plugins/MinikuraVelocity/src/main/kotlin/cafe/kirameki/minikuraVelocity/listeners/ProxyTransferHandler.kt b/plugins/MinikuraVelocity/src/main/kotlin/cafe/kirameki/minikuraVelocity/listeners/ProxyTransferHandler.kt index 462aead..7c3e177 100644 --- a/plugins/MinikuraVelocity/src/main/kotlin/cafe/kirameki/minikuraVelocity/listeners/ProxyTransferHandler.kt +++ b/plugins/MinikuraVelocity/src/main/kotlin/cafe/kirameki/minikuraVelocity/listeners/ProxyTransferHandler.kt @@ -5,19 +5,23 @@ import com.auth0.jwt.JWT import com.auth0.jwt.algorithms.Algorithm import com.imaginarycode.minecraft.redisbungee.RedisBungeeAPI import com.velocitypowered.api.event.Subscribe +import com.velocitypowered.api.event.connection.PreLoginEvent import com.velocitypowered.api.event.player.CookieReceiveEvent import com.velocitypowered.api.event.player.PlayerChooseInitialServerEvent import com.velocitypowered.api.proxy.server.RegisteredServer import net.kyori.adventure.key.Key +import net.kyori.adventure.text.Component import org.slf4j.Logger import java.util.* import java.util.concurrent.CompletableFuture import java.util.concurrent.TimeUnit import java.util.concurrent.TimeoutException +import java.util.concurrent.atomic.AtomicBoolean class ProxyTransferHandler( private val servers: Map, - private val logger: Logger + private val logger: Logger, + private val acceptingTransfer: AtomicBoolean ) { private val cookieFutures = mutableMapOf>() private var jwtSecret = System.getenv("MINIKURA_JWT_SECRET") ?: "secret" @@ -25,6 +29,13 @@ class ProxyTransferHandler( private val jwtVerifier = JWT.require(jwtAlgorithm).withIssuer("minikura").build() private val redisBungeeApi = RedisBungeeAPI.getRedisBungeeApi() + @Subscribe + fun onPreLogin(event: PreLoginEvent) { + if (!acceptingTransfer.get()) { + event.result = PreLoginEvent.PreLoginComponentResult.denied(Component.text("Proxy are currently not accepting new connections.")) + } + } + @Subscribe fun onCookieReceiveEvent(event: CookieReceiveEvent) { // TODO: Fix don't pass the cookie to backend, it kicks with invalid packet for some reason